"""GEOPipeline - GEO 端到端工作流编排

实现检测→分析→优化→追踪的 DAG Pipeline，
基于 Orchestrator 的多 Agent 协作模式。
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from dataclasses import dataclass, field
from typing import Any

from agentkit.core.protocol import TaskMessage
from agentkit.core.shared_workspace import SharedWorkspace
from agentkit.orchestrator.compensation import SagaOrchestrator
from agentkit.orchestrator.retry import StepRetryPolicy, execute_with_retry
from agentkit.skills.registry import SkillRegistry

logger = logging.getLogger(__name__)


@dataclass
class PipelineStep:
    """Pipeline 步骤定义"""

    name: str
    skill: str
    input_mapping: dict[str, str] = field(default_factory=dict)
    depends_on: list[str] = field(default_factory=list)
    condition: str | None = None
    parallel_with: list[str] = field(default_factory=list)
    compensate: str | None = None
    retry_policy: StepRetryPolicy | None = None


@dataclass
class PipelineStepResult:
    """步骤执行结果"""

    step_name: str
    skill: str
    status: str  # "success", "failed", "skipped"
    output: dict[str, Any] | None = None
    error: str | None = None
    duration_ms: float = 0


@dataclass
class PipelineResult:
    """Pipeline 执行结果"""

    pipeline_name: str
    execution_id: str
    steps: list[PipelineStepResult]
    final_output: dict[str, Any] | None
    success: bool
    total_duration_ms: float


class GEOPipeline:
    """GEO 端到端工作流编排

    支持：
    - YAML 配置驱动的 Pipeline 定义
    - DAG 依赖关系（depends_on）
    - 并行执行无依赖的步骤
    - 步骤间数据通过 SharedWorkspace 传递
    - 条件跳过步骤

    使用方式：
        pipeline = GEOPipeline.from_config(config, skill_registry, agent_pool)
        result = await pipeline.execute(input_data)
    """

    def __init__(
        self,
        name: str,
        steps: list[PipelineStep],
        skill_registry: SkillRegistry | None = None,
        agent_pool: Any = None,
        workspace: SharedWorkspace | None = None,
    ):
        self.name = name
        self._steps = steps
        self._skill_registry = skill_registry
        self._agent_pool = agent_pool
        self._workspace = workspace or SharedWorkspace()
        self._step_map = {s.name: s for s in steps}

    @classmethod
    def from_config(
        cls,
        config: dict[str, Any],
        skill_registry: SkillRegistry | None = None,
        agent_pool: Any = None,
        workspace: SharedWorkspace | None = None,
    ) -> GEOPipeline:
        """从 YAML 配置创建 Pipeline

        配置格式：
            name: geo_full_pipeline
            steps:
              - name: detect
                skill: citation_detector
                input_mapping: {brand: $.input.brand}
              - name: analyze
                skill: competitor_analyzer
                depends_on: [detect]
        """
        steps = []
        for step_conf in config.get("steps", []):
            retry_policy = None
            retry_conf = step_conf.get("retry_policy")
            if retry_conf:
                retry_policy = StepRetryPolicy(**retry_conf)

            step = PipelineStep(
                name=step_conf["name"],
                skill=step_conf["skill"],
                input_mapping=step_conf.get("input_mapping", {}),
                depends_on=step_conf.get("depends_on", []),
                condition=step_conf.get("condition"),
                parallel_with=step_conf.get("parallel_with", []),
                compensate=step_conf.get("compensate"),
                retry_policy=retry_policy,
            )
            steps.append(step)

        return cls(
            name=config.get("name", "geo_pipeline"),
            steps=steps,
            skill_registry=skill_registry,
            agent_pool=agent_pool,
            workspace=workspace,
        )

    async def execute(self, input_data: dict[str, Any]) -> PipelineResult:
        """执行 Pipeline

        Args:
            input_data: 初始输入数据

        Returns:
            PipelineResult: 包含各步骤结果和最终输出
        """
        import time

        start_time = time.monotonic()
        execution_id = str(uuid.uuid4())[:8]
        step_results: list[PipelineStepResult] = []
        step_outputs: dict[str, dict[str, Any]] = {}

        # Store initial input in workspace
        await self._workspace.write(
            f"pipeline:{execution_id}:input",
            input_data,
            agent_id="pipeline",
        )

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

        # Build execution order (topological sort)
        execution_groups = self._build_execution_groups()

        pipeline_failed = False
        for group in execution_groups:
            # Execute group in parallel
            tasks = []
            for step_name in group:
                step = self._step_map[step_name]
                tasks.append(self._execute_step(step, input_data, step_outputs, execution_id, saga))

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

            for step_name, result in zip(group, results):
                if isinstance(result, Exception):
                    step_result = PipelineStepResult(
                        step_name=step_name,
                        skill=self._step_map[step_name].skill,
                        status="failed",
                        error=str(result),
                    )
                else:
                    step_result = result

                step_results.append(step_result)
                if step_result.status == "success" and step_result.output:
                    step_outputs[step_name] = step_result.output

                # On failure, trigger Saga compensation
                if step_result.status == "failed":
                    pipeline_failed = True
                    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]}"
                            )
                    break

            if pipeline_failed:
                break

        # Build final output
        final_output = self._build_final_output(step_outputs, input_data)

        duration_ms = (time.monotonic() - start_time) * 1000
        success = all(r.status in ("success", "skipped") for r in step_results)

        return PipelineResult(
            pipeline_name=self.name,
            execution_id=execution_id,
            steps=step_results,
            final_output=final_output,
            success=success,
            total_duration_ms=duration_ms,
        )

    async def _execute_step(
        self,
        step: PipelineStep,
        input_data: dict[str, Any],
        step_outputs: dict[str, dict[str, Any]],
        execution_id: str,
        saga: SagaOrchestrator,
    ) -> PipelineStepResult:
        """执行单个 Pipeline 步骤"""
        import time

        start_time = time.monotonic()

        # Check condition
        if step.condition and not self._evaluate_condition(step.condition, input_data, step_outputs):
            return PipelineStepResult(
                step_name=step.name,
                skill=step.skill,
                status="skipped",
            )

        # Build step input from mapping
        step_input = self._map_input(step, input_data, step_outputs)

        # Execute skill (with retry if configured)
        try:
            if step.retry_policy is not None:
                output = await execute_with_retry(
                    func=lambda: self._execute_skill(step.skill, step_input),
                    retry_policy=step.retry_policy,
                    step_name=step.name,
                )
            else:
                output = await self._execute_skill(step.skill, step_input)

            duration_ms = (time.monotonic() - start_time) * 1000

            # Store result in workspace
            await self._workspace.write(
                f"pipeline:{execution_id}:step:{step.name}",
                output,
                agent_id=step.skill,
            )

            # Record completed step for Saga compensation
            saga.record_completed(
                step_name=step.name,
                result=output,
                compensate_action=step.compensate,
            )

            return PipelineStepResult(
                step_name=step.name,
                skill=step.skill,
                status="success",
                output=output,
                duration_ms=duration_ms,
            )
        except Exception as e:
            duration_ms = (time.monotonic() - start_time) * 1000
            logger.error(f"Pipeline step '{step.name}' failed: {e}")
            return PipelineStepResult(
                step_name=step.name,
                skill=step.skill,
                status="failed",
                error=str(e),
                duration_ms=duration_ms,
            )

    async def _execute_skill(
        self, skill_name: str, input_data: dict[str, Any]
    ) -> dict[str, Any]:
        """执行 Skill"""
        if self._agent_pool:
            agent = self._agent_pool.get_agent(skill_name)
            if agent:
                from datetime import datetime, timezone

                task = TaskMessage(
                    task_id=f"pipeline-{skill_name}",
                    agent_name=skill_name,
                    task_type=skill_name,
                    priority=0,
                    input_data=input_data,
                    callback_url=None,
                    created_at=datetime.now(timezone.utc),
                )
                result = await agent.execute(task)
                return result.output_data if hasattr(result, "output_data") else result

        if self._skill_registry:
            skill = self._skill_registry.get(skill_name)
            from agentkit.core.config_driven import ConfigDrivenAgent
            from datetime import datetime, timezone

            agent = ConfigDrivenAgent(config=skill.config)
            task = TaskMessage(
                task_id=f"pipeline-{skill_name}",
                agent_name=skill_name,
                task_type=skill.config.agent_type,
                priority=0,
                input_data=input_data,
                callback_url=None,
                created_at=datetime.now(timezone.utc),
            )
            return await agent.handle_task(task)

        raise ValueError(f"Skill '{skill_name}' not found: no agent_pool or skill_registry")

    def _build_execution_groups(self) -> list[list[str]]:
        """构建并行执行组（拓扑排序）"""
        completed: set[str] = set()
        groups: list[list[str]] = []
        remaining = set(s.name for s in self._steps)

        while remaining:
            ready = []
            for name in remaining:
                step = self._step_map[name]
                if all(dep in completed for dep in step.depends_on):
                    ready.append(name)

            if not ready:
                # Circular dependency — force remaining into one group
                groups.append(list(remaining))
                break

            groups.append(ready)
            for name in ready:
                completed.add(name)
                remaining.discard(name)

        return groups

    def _map_input(
        self,
        step: PipelineStep,
        input_data: dict[str, Any],
        step_outputs: dict[str, dict[str, Any]],
    ) -> dict[str, Any]:
        """根据 input_mapping 构建步骤输入

        映射格式: {"target_key": "source_path"}
        source_path 支持:
        - $.input.xxx — 初始输入
        - $.steps.step_name.output.xxx — 步骤输出
        """
        if not step.input_mapping:
            # Default: merge all dependency outputs + original input
            merged = dict(input_data)
            for dep in step.depends_on:
                if dep in step_outputs:
                    merged.update(step_outputs[dep])
            return merged

        mapped: dict[str, Any] = {}
        for target_key, source_path in step.input_mapping.items():
            value = self._resolve_mapping_path(source_path, input_data, step_outputs)
            if value is not None:
                mapped[target_key] = value

        return mapped

    @staticmethod
    def _resolve_mapping_path(
        path: str,
        input_data: dict[str, Any],
        step_outputs: dict[str, dict[str, Any]],
    ) -> Any:
        """解析映射路径"""
        if path.startswith("$.input."):
            key = path[len("$.input."):]
            return input_data.get(key)
        elif path.startswith("$.steps."):
            # $.steps.step_name or $.steps.step_name.output.field
            rest = path[len("$.steps."):]
            parts = rest.split(".", 2)
            step_name = parts[0]
            if step_name not in step_outputs:
                return None
            if len(parts) == 1:
                # $.steps.step_name — return whole output
                return step_outputs[step_name]
            if len(parts) >= 2 and parts[1] == "output":
                if len(parts) >= 3:
                    return step_outputs[step_name].get(parts[2])
                return step_outputs[step_name]
            # $.steps.step_name.field (without .output)
            return step_outputs[step_name].get(parts[1])
        return None

    def _evaluate_condition(
        self, condition: str, input_data: dict[str, Any], step_outputs: dict[str, Any]
    ) -> bool:
        """评估条件表达式"""
        import re

        try:
            eq_match = re.match(r'^([\w.]+)\s*==\s*(.+)$', condition.strip())
            if eq_match:
                path = eq_match.group(1)
                value = eq_match.group(2).strip().strip("'\"")
                actual = self._resolve_mapping_path(f"$.{path}", input_data, step_outputs)
                return str(actual) == value
        except (ValueError, TypeError) as e:
            logger.warning(f"Condition evaluation failed for '{condition}': {e}")
            return False
        return True

    def _build_final_output(
        self,
        step_outputs: dict[str, dict[str, Any]],
        input_data: dict[str, Any],
    ) -> dict[str, Any]:
        """构建最终输出"""
        final = {"input": input_data}
        for step_name, output in step_outputs.items():
            final[step_name] = output
        return final


# GEO Pipeline 默认步骤补偿定义
GEO_PIPELINE_COMPENSATIONS: dict[str, str | None] = {
    "detect": None,                    # 只读操作，无需补偿
    "analyze_competitor": None,        # 只读操作，无需补偿
    "optimize": "revert_optimization", # 需要回滚优化变更
    "schema": None,                    # 幂等操作，无需补偿
    "monitor": None,                   # 只读操作，无需补偿
}


def create_geo_pipeline_steps() -> list[PipelineStep]:
    """创建 GEO Pipeline 默认步骤（含补偿定义）"""
    steps = [
        PipelineStep(
            name="detect",
            skill="citation_detector",
            input_mapping={"brand": "$.input.brand"},
            compensate=GEO_PIPELINE_COMPENSATIONS["detect"],
        ),
        PipelineStep(
            name="analyze_competitor",
            skill="competitor_analyzer",
            depends_on=["detect"],
            input_mapping={"brand": "$.input.brand"},
            compensate=GEO_PIPELINE_COMPENSATIONS["analyze_competitor"],
        ),
        PipelineStep(
            name="optimize",
            skill="content_optimizer",
            depends_on=["analyze_competitor"],
            input_mapping={"brand": "$.input.brand"},
            compensate=GEO_PIPELINE_COMPENSATIONS["optimize"],
        ),
        PipelineStep(
            name="schema",
            skill="schema_generator",
            depends_on=["optimize"],
            input_mapping={"brand": "$.input.brand"},
            compensate=GEO_PIPELINE_COMPENSATIONS["schema"],
        ),
        PipelineStep(
            name="monitor",
            skill="citation_monitor",
            depends_on=["schema"],
            input_mapping={"brand": "$.input.brand"},
            compensate=GEO_PIPELINE_COMPENSATIONS["monitor"],
        ),
    ]
    return steps
