"""Pipeline Engine - DAG + 并行执行 + 步骤重试 + Saga 补偿"""

import asyncio
import logging
from collections import defaultdict
from datetime import datetime, timezone
from typing import Any

from agentkit.orchestrator.compensation import SagaOrchestrator
from agentkit.orchestrator.pipeline_schema import (
    AdaptiveConfig,
    Pipeline,
    PipelineResult,
    PipelineStage,
    ReflectionReport,
    StageResult,
    StageStatus,
)
from agentkit.orchestrator.reflection import PipelineReflector, PipelineReplanner
from agentkit.orchestrator.retry import StepRetryPolicy, execute_with_retry

logger = logging.getLogger(__name__)


class PipelineEngine:
    """Pipeline 执行引擎

    支持：
    - DAG 拓扑排序
    - 同层并行执行（asyncio.gather）
    - 变量解析
    - 条件执行
    - 步骤级指数退避重试（StepRetryPolicy）
    - Saga 补偿（LIFO 回滚已完成步骤）
    - 状态持久化（可选）
    """

    def __init__(self, dispatcher: Any = None, state_manager: Any = None, llm_gateway: Any = None):
        self._dispatcher = dispatcher
        self._state_manager = state_manager
        self._llm_gateway = llm_gateway

    async def execute(
        self,
        pipeline: Pipeline,
        context: dict[str, Any] | None = None,
        adaptive_config: AdaptiveConfig | None = None,
    ) -> PipelineResult:
        """执行 Pipeline

        Args:
            pipeline: Pipeline 定义
            context: 运行时上下文变量
            adaptive_config: 自适应配置，启用反思-重规划闭环
        """
        # First execution
        result = await self._execute_pipeline(pipeline, context)

        # If failed and adaptive is enabled, enter reflection-replanning loop
        if result.status == StageStatus.FAILED and adaptive_config and adaptive_config.enabled:
            result = await self._adaptive_loop(pipeline, context, result, adaptive_config)

        return result

    async def _adaptive_loop(
        self,
        pipeline: Pipeline,
        context: dict[str, Any] | None,
        failed_result: PipelineResult,
        adaptive_config: AdaptiveConfig,
    ) -> PipelineResult:
        """反思-重规划闭环：分析失败原因 → 修正 Pipeline → 重新执行。"""
        reflector = PipelineReflector(llm_gateway=self._llm_gateway)
        replanner = PipelineReplanner(llm_gateway=self._llm_gateway)

        current_pipeline = pipeline
        current_result = failed_result
        reflections: list[ReflectionReport] = []

        for reflection_num in range(1, adaptive_config.max_reflections + 1):
            # Reflect
            report = await reflector.reflect(current_pipeline, current_result, reflection_num)
            reflections.append(report)
            logger.info(
                f"Pipeline reflection #{reflection_num}: "
                f"failure_type={report.failure_type}, "
                f"root_cause={report.root_cause}"
            )

            # Replan
            new_pipeline = await replanner.replan(current_pipeline, current_result, report)
            logger.info(f"Pipeline replanned: {new_pipeline.name} ({len(new_pipeline.stages)} stages)")

            # Re-execute
            current_result = await self._execute_pipeline(new_pipeline, context)
            current_pipeline = new_pipeline

            # Record reflection in metadata
            current_result.metadata["reflections"] = [
                r.model_dump() for r in reflections
            ]

            if current_result.status == StageStatus.COMPLETED:
                logger.info(f"Pipeline succeeded after {reflection_num} reflection(s)")
                return current_result

        # Exhausted reflections
        logger.warning(
            f"Pipeline failed after {adaptive_config.max_reflections} reflection(s)"
        )
        current_result.metadata["reflections"] = [
            r.model_dump() for r in reflections
        ]
        return current_result

    async def _execute_pipeline(
        self,
        pipeline: Pipeline,
        context: dict[str, Any] | None = None,
    ) -> PipelineResult:
        """执行 Pipeline 的核心逻辑（不含反思-重规划）。"""
        result = PipelineResult(pipeline_name=pipeline.name)
        result.variables = {**pipeline.variables, **(context or {})}

        # 拓扑排序 + 按依赖层级分组
        try:
            level_groups = self._topological_group(pipeline.stages)
        except ValueError as e:
            result.status = StageStatus.FAILED
            result.error_message = str(e)
            return result

        # Create execution state if state_manager is configured
        execution_id: str | None = None
        if self._state_manager is not None:
            try:
                step_names = [s.name for s in pipeline.stages]
                execution_id = await self._state_manager.create_execution(
                    pipeline_name=pipeline.name,
                    steps=step_names,
                    input_data=context,
                )
            except Exception as exc:
                logger.warning(f"Failed to create execution state: {exc}")

        # Create Saga orchestrator for compensation tracking
        saga = SagaOrchestrator()

        # 逐层执行
        for level, stages in enumerate(level_groups):
            logger.info(f"Pipeline '{pipeline.name}' executing level {level} with {len(stages)} stage(s)")

            # 并行执行同层 stages
            tasks = []
            for stage in stages:
                tasks.append(self._execute_stage(stage, result, saga))

            stage_results = await asyncio.gather(*tasks, return_exceptions=True)

            # 处理结果
            for stage, sr in zip(stages, stage_results):
                if isinstance(sr, Exception):
                    sr = StageResult(
                        stage_name=stage.name,
                        status=StageStatus.FAILED,
                        error_message=str(sr),
                    )
                result.stage_results[stage.name] = sr

                # Update step state
                if self._state_manager is not None and execution_id is not None:
                    try:
                        step_status = "completed" if sr.status == StageStatus.COMPLETED else sr.status.value
                        step_output = sr.output_data if hasattr(sr, 'output_data') else None
                        step_error = sr.error_message if hasattr(sr, 'error_message') else None
                        await self._state_manager.update_step(
                            execution_id=execution_id,
                            step_name=stage.name,
                            status=step_status,
                            output=step_output,
                            error=step_error,
                        )
                    except Exception as exc:
                        logger.warning(f"Failed to update step state: {exc}")

                # 收集输出变量
                if sr.output_data and isinstance(sr, dict):
                    pass
                elif hasattr(sr, 'output_data') and sr.output_data:
                    for output_key in stage.outputs:
                        if output_key in sr.output_data:
                            result.variables[output_key] = sr.output_data[output_key]

                # 检查是否需要中止
                if hasattr(sr, 'status') and sr.status == StageStatus.FAILED:
                    if not stage.continue_on_failure:
                        # Execute Saga compensation for completed steps
                        compensation_results = await saga.compensate()
                        if compensation_results:
                            failed_compensations = [
                                cr for cr in compensation_results if not cr.success and cr.error != "no_compensation_needed"
                            ]
                            if failed_compensations:
                                logger.warning(
                                    f"Compensation had {len(failed_compensations)} failures: "
                                    f"{[c.step_name for c in failed_compensations]}"
                                )

                        result.status = StageStatus.FAILED
                        result.error_message = f"Stage '{stage.name}' failed"
                        # Fail execution state
                        if self._state_manager is not None and execution_id is not None:
                            try:
                                await self._state_manager.fail_execution(
                                    execution_id=execution_id,
                                    step_name=stage.name,
                                    error=result.error_message,
                                )
                            except Exception as exc:
                                logger.warning(f"Failed to persist failure state: {exc}")
                        return result

        result.status = StageStatus.COMPLETED

        # Complete execution state
        if self._state_manager is not None and execution_id is not None:
            try:
                final_output = {
                    name: sr.output_data
                    for name, sr in result.stage_results.items()
                    if sr.output_data is not None
                }
                await self._state_manager.complete_execution(
                    execution_id=execution_id,
                    final_output=final_output,
                )
            except Exception as exc:
                logger.warning(f"Failed to persist completion state: {exc}")

        return result

    async def _execute_stage(
        self,
        stage: PipelineStage,
        pipeline_result: PipelineResult,
        saga: SagaOrchestrator,
    ) -> StageResult:
        """执行单个 stage"""
        started_at = datetime.now(timezone.utc).isoformat()

        # 条件检查
        if stage.condition and not self._evaluate_condition(stage.condition, pipeline_result.variables):
            return StageResult(
                stage_name=stage.name,
                status=StageStatus.SKIPPED,
                started_at=started_at,
                completed_at=datetime.now(timezone.utc).isoformat(),
            )

        # 解析输入变量
        resolved_inputs = self._resolve_variables(stage.inputs, pipeline_result.variables)

        # 执行
        if self._dispatcher is None:
            # Dry-run 模式
            result = StageResult(
                stage_name=stage.name,
                status=StageStatus.COMPLETED,
                output_data={"dry_run": True, "inputs": resolved_inputs},
                started_at=started_at,
                completed_at=datetime.now(timezone.utc).isoformat(),
            )
            # Record completed step for Saga compensation
            saga.record_completed(
                step_name=stage.name,
                result=result.output_data,
                compensate_action=stage.compensate,
            )
            return result

        # 通过 Dispatcher 分发任务
        from agentkit.core.protocol import TaskMessage
        import uuid

        task = TaskMessage(
            task_id=str(uuid.uuid4()),
            agent_name=stage.agent,
            task_type=stage.action,
            priority=0,
            input_data=resolved_inputs,
            callback_url=None,
            created_at=datetime.now(timezone.utc),
            timeout_seconds=stage.timeout_seconds,
        )

        async def _dispatch_and_wait() -> StageResult:
            """Dispatch task and wait for result"""
            await self._dispatcher.dispatch(task)

            # 等待结果
            for _ in range(stage.timeout_seconds):
                status = await self._dispatcher.get_task_status(task.task_id)
                if status["status"] in ("completed", "failed", "cancelled"):
                    return StageResult(
                        stage_name=stage.name,
                        status=StageStatus.COMPLETED if status["status"] == "completed" else StageStatus.FAILED,
                        output_data=status.get("output_data"),
                        error_message=status.get("error_message"),
                        started_at=started_at,
                        completed_at=datetime.now(timezone.utc).isoformat(),
                    )
                await asyncio.sleep(1)

            return StageResult(
                stage_name=stage.name,
                status=StageStatus.FAILED,
                error_message=f"Timeout after {stage.timeout_seconds}s",
                started_at=started_at,
                completed_at=datetime.now(timezone.utc).isoformat(),
            )

        try:
            # Execute with retry if retry_policy is configured
            sr = await execute_with_retry(
                func=_dispatch_and_wait,
                retry_policy=stage.retry_policy,
                step_name=stage.name,
            )

            # Record completed step for Saga compensation on success
            if sr.status == StageStatus.COMPLETED:
                saga.record_completed(
                    step_name=stage.name,
                    result=sr.output_data,
                    compensate_action=stage.compensate,
                )

            return sr

        except Exception as e:
            return StageResult(
                stage_name=stage.name,
                status=StageStatus.FAILED,
                error_message=str(e),
                started_at=started_at,
                completed_at=datetime.now(timezone.utc).isoformat(),
            )

    @staticmethod
    def _topological_group(stages: list[PipelineStage]) -> list[list[PipelineStage]]:
        """拓扑排序 + 按依赖层级分组"""
        stage_map = {s.name: s for s in stages}
        in_degree = defaultdict(int)
        dependents = defaultdict(list)

        for s in stages:
            if s.name not in in_degree:
                in_degree[s.name] = 0
            for dep in s.depends_on:
                if dep not in stage_map:
                    raise ValueError(f"Stage '{s.name}' depends on unknown stage '{dep}'")
                in_degree[s.name] += 1
                dependents[dep].append(s.name)

        levels = []
        remaining = set(in_degree.keys())

        while remaining:
            # 找到入度为 0 的节点
            current_level = [name for name in remaining if in_degree[name] == 0]
            if not current_level:
                raise ValueError("Circular dependency detected in pipeline")

            levels.append([stage_map[name] for name in current_level])

            for name in current_level:
                remaining.remove(name)
                for dep in dependents[name]:
                    in_degree[dep] -= 1

        return levels

    @staticmethod
    def _resolve_variables(template: dict, context: dict) -> dict:
        """解析 ${var.path} 变量引用"""
        resolved = {}
        for key, value in template.items():
            if isinstance(value, str) and value.startswith("${") and value.endswith("}"):
                var_path = value[2:-1]
                resolved[key] = PipelineEngine._get_nested(context, var_path)
            else:
                resolved[key] = value
        return resolved

    @staticmethod
    def _get_nested(data: dict, path: str) -> Any:
        keys = path.split(".")
        current = data
        for key in keys:
            if isinstance(current, dict):
                current = current.get(key)
            else:
                return None
        return current

    @staticmethod
    def _evaluate_condition(condition: str, variables: dict) -> bool:
        """简单条件评估"""
        if "==" in condition:
            parts = condition.split("==", 1)
            left = variables.get(parts[0].strip(), parts[0].strip())
            right = parts[1].strip().strip("'\"")
            return str(left) == right
        elif "!=" in condition:
            parts = condition.split("!=", 1)
            left = variables.get(parts[0].strip(), parts[0].strip())
            right = parts[1].strip().strip("'\"")
            return str(left) != right
        else:
            return bool(variables.get(condition))
