"""Orchestrator - 多 Agent 协作编排器

实现 Orchestrator-Worker 模式：中央编排器协调多 Agent 并行/串行执行。
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any

from agentkit.core.protocol import TaskMessage, TaskResult, TaskStatus
from agentkit.core.shared_workspace import SharedWorkspace

if TYPE_CHECKING:
    from agentkit.core.goal_planner import GoalPlanner
    from agentkit.core.plan_executor import PlanExecutor
    from agentkit.core.plan_checker import PlanChecker

logger = logging.getLogger(__name__)


class AgentRole(str, Enum):
    """Agent 角色枚举"""

    ORCHESTRATOR = "orchestrator"
    WORKER = "worker"
    REVIEWER = "reviewer"


class SubTaskStatus(str, Enum):
    """子任务状态"""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class SubTask:
    """子任务定义"""

    task_id: str
    parent_task_id: str
    assigned_agent: str
    task_type: str
    input_data: dict[str, Any]
    status: SubTaskStatus = SubTaskStatus.PENDING
    result: dict[str, Any] | None = None
    error: str | None = None
    depends_on: list[str] = field(default_factory=list)


@dataclass
class OrchestrationPlan:
    """编排计划"""

    plan_id: str
    parent_task_id: str
    subtasks: list[SubTask]
    parallel_groups: list[list[str]]  # 每组内的子任务可并行执行


@dataclass
class OrchestrationResult:
    """编排结果"""

    plan_id: str
    parent_task_id: str
    subtask_results: dict[str, dict[str, Any]]
    aggregated_result: dict[str, Any]
    status: TaskStatus
    total_duration_ms: float
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class OrchestratorConfig:
    """Orchestrator 配置"""

    adaptive: bool = False
    max_iterations: int = 3
    quality_threshold: float = 0.7


class Orchestrator:
    """多 Agent 协作编排器

    Orchestrator-Worker 模式：
    1. 接收复杂任务
    2. LLM 驱动分解为子任务
    3. 基于 Skill 能力匹配子任务到 Worker Agent
    4. 并行/串行执行子任务
    5. 汇总结果，生成最终输出

    使用方式：
        orchestrator = Orchestrator(agent_pool=pool, workspace=workspace)
        result = await orchestrator.execute(task_message)
    """

    def __init__(
        self,
        agent_pool: Any,
        workspace: SharedWorkspace | None = None,
        llm_gateway: Any = None,
        max_parallel: int = 5,
        subtask_timeout: float = 300.0,
        goal_planner: GoalPlanner | None = None,
        plan_executor: PlanExecutor | None = None,
        plan_checker: PlanChecker | None = None,
        config: OrchestratorConfig | None = None,
        message_bus: Any = None,
    ):
        """
        Args:
            agent_pool: AgentPool 实例
            workspace: 共享工作空间
            llm_gateway: LLM Gateway，用于任务分解
            max_parallel: 最大并行子任务数
            subtask_timeout: 子任务超时时间（秒）
            goal_planner: GoalPlanner 实例，用于结构化目标分解（可选）
            plan_executor: PlanExecutor 实例，用于执行 ExecutionPlan（可选）
            plan_checker: PlanChecker 实例，用于检查和复盘（可选）
            config: Orchestrator 配置，包含自适应参数
            message_bus: MessageBus 实例，用于 Agent 间通信
        """
        self._agent_pool = agent_pool
        self._workspace = workspace or SharedWorkspace()
        self._llm_gateway = llm_gateway
        self._max_parallel = max_parallel
        self._subtask_timeout = subtask_timeout
        self._goal_planner = goal_planner
        self._plan_executor = plan_executor
        self._plan_checker = plan_checker
        self._config = config or OrchestratorConfig()
        self._message_bus = message_bus

    async def execute(self, task: TaskMessage) -> OrchestrationResult:
        """执行编排任务

        Args:
            task: 原始任务消息

        Returns:
            OrchestrationResult: 编排结果
        """
        import time

        start_time = time.monotonic()

        # 1. Decompose task into subtasks
        plan = await self._decompose_task(task)

        if not plan.subtasks:
            return OrchestrationResult(
                plan_id=plan.plan_id,
                parent_task_id=task.task_id,
                subtask_results={},
                aggregated_result={"error": "Failed to decompose task"},
                status=TaskStatus.FAILED,
                total_duration_ms=0,
            )

        # 2. Store plan in workspace
        await self._workspace.write(
            f"plan:{plan.plan_id}",
            {"task_id": task.task_id, "subtask_count": len(plan.subtasks)},
            agent_id="orchestrator",
        )

        # 3. Execute subtasks
        subtask_results = await self._execute_plan(plan, task)

        # 4. Aggregate results
        aggregated = await self._aggregate_results(plan, subtask_results, task)

        # 5. Determine overall status
        failed_count = sum(
            1 for r in subtask_results.values() if r.get("status") == "failed"
        )
        if failed_count == len(plan.subtasks):
            status = TaskStatus.FAILED
        elif failed_count > 0:
            status = TaskStatus.PARTIALLY_COMPLETED
        else:
            status = TaskStatus.COMPLETED

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

        return OrchestrationResult(
            plan_id=plan.plan_id,
            parent_task_id=task.task_id,
            subtask_results=subtask_results,
            aggregated_result=aggregated,
            status=status,
            total_duration_ms=duration_ms,
        )

    async def _decompose_task(self, task: TaskMessage) -> OrchestrationPlan:
        """将复杂任务分解为子任务"""
        plan_id = str(uuid.uuid4())[:8]

        # If GoalPlanner available, use it for structured decomposition
        if self._goal_planner:
            try:
                execution_plan = await self._goal_planner.generate_plan(
                    goal=str(task.input_data),
                    context={"task_type": task.task_type, "agent_name": task.agent_name},
                    available_skills=self._get_available_skill_names(),
                )
                subtasks = self._convert_execution_plan_to_subtasks(
                    execution_plan, task.task_id, task.agent_name, task.task_type, task.input_data,
                )
                if subtasks:
                    parallel_groups = self._build_parallel_groups(subtasks)
                    return OrchestrationPlan(
                        plan_id=plan_id,
                        parent_task_id=task.task_id,
                        subtasks=subtasks,
                        parallel_groups=parallel_groups,
                    )
            except Exception as e:
                logger.warning(f"GoalPlanner decomposition failed, falling back: {e}")

        # If LLM gateway available, use it for decomposition
        if self._llm_gateway:
            try:
                subtasks = await self._llm_decompose(task)
                if subtasks:
                    parallel_groups = self._build_parallel_groups(subtasks)
                    return OrchestrationPlan(
                        plan_id=plan_id,
                        parent_task_id=task.task_id,
                        subtasks=subtasks,
                        parallel_groups=parallel_groups,
                    )
            except Exception as e:
                logger.warning(f"LLM decomposition failed, falling back to simple: {e}")

        # Fallback: single subtask = original task
        subtask = SubTask(
            task_id=f"{plan_id}-0",
            parent_task_id=task.task_id,
            assigned_agent=task.agent_name,
            task_type=task.task_type,
            input_data=task.input_data,
        )
        return OrchestrationPlan(
            plan_id=plan_id,
            parent_task_id=task.task_id,
            subtasks=[subtask],
            parallel_groups=[[subtask.task_id]],
        )

    async def _llm_decompose(self, task: TaskMessage) -> list[SubTask]:
        """使用 LLM 分解任务"""
        # Get available agents and their capabilities
        agents_info = self._agent_pool.list_agents()
        agent_descriptions = "\n".join(
            f"- {a['name']} ({a['agent_type']}): {a.get('description', 'No description')}"
            for a in agents_info
        )

        prompt = (
            f"Decompose the following task into subtasks that can be assigned to available agents.\n\n"
            f"Task: {task.input_data}\n"
            f"Task Type: {task.task_type}\n\n"
            f"Available Agents:\n{agent_descriptions}\n\n"
            'Respond ONLY with a JSON array: [{"agent_name": "...", "task_type": "...", '
            '"input_data": {...}, "depends_on": []}]\n'
            "The depends_on field lists task indices (0-based) that must complete first.\n"
            "Do not include any other text."
        )

        import json

        response = await self._llm_gateway.chat(
            messages=[{"role": "user", "content": prompt}],
            model="default",
        )

        try:
            subtask_defs = json.loads(response.content)
            if not isinstance(subtask_defs, list):
                return []

            subtasks = []
            for i, defn in enumerate(subtask_defs):
                depends_on = [
                    f"task-{i}" for i in defn.get("depends_on", [])
                    if isinstance(i, int) and 0 <= i < len(subtask_defs)
                ]
                subtasks.append(SubTask(
                    task_id=f"task-{i}",
                    parent_task_id=task.task_id,
                    assigned_agent=defn.get("agent_name", task.agent_name),
                    task_type=defn.get("task_type", task.task_type),
                    input_data=defn.get("input_data", {}),
                    depends_on=depends_on,
                ))
            return subtasks
        except (json.JSONDecodeError, KeyError) as e:
            logger.warning(f"Failed to parse LLM decomposition: {e}")
            return []

    def _build_parallel_groups(self, subtasks: list[SubTask]) -> list[list[str]]:
        """构建并行执行组

        基于依赖关系拓扑排序，无依赖的子任务分到同一组并行执行。
        """
        # Build dependency graph
        task_map = {st.task_id: st for st in subtasks}
        completed: set[str] = set()
        groups: list[list[str]] = []

        remaining = set(st.task_id for st in subtasks)

        while remaining:
            # Find tasks with all dependencies satisfied
            ready = []
            for tid in remaining:
                task = task_map[tid]
                if all(dep in completed for dep in task.depends_on):
                    ready.append(tid)

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

            # Limit group size
            group = ready[:self._max_parallel]
            groups.append(group)
            for tid in group:
                completed.add(tid)
                remaining.discard(tid)

        return groups

    async def _execute_plan(
        self, plan: OrchestrationPlan, original_task: TaskMessage
    ) -> dict[str, dict[str, Any]]:
        """执行编排计划"""
        subtask_results: dict[str, dict[str, Any]] = {}
        task_map = {st.task_id: st for st in plan.subtasks}

        for group in plan.parallel_groups:
            # Execute group in parallel
            tasks = []
            for task_id in group:
                subtask = task_map[task_id]
                # Inject results from dependencies
                enriched_input = self._inject_dependency_results(
                    subtask, subtask_results
                )
                tasks.append(self._execute_subtask(subtask, enriched_input, original_task))

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

            for task_id, result in zip(group, results):
                if isinstance(result, Exception):
                    subtask_results[task_id] = {
                        "status": "failed",
                        "error": str(result),
                    }
                else:
                    subtask_results[task_id] = result

        return subtask_results

    async def _execute_subtask(
        self,
        subtask: SubTask,
        input_data: dict[str, Any],
        original_task: TaskMessage,
    ) -> dict[str, Any]:
        """执行单个子任务"""
        agent = self._agent_pool.get_agent(subtask.assigned_agent)
        if agent is None:
            return {"status": "failed", "error": f"Agent '{subtask.assigned_agent}' not found"}

        sub_task_msg = TaskMessage(
            task_id=subtask.task_id,
            agent_name=subtask.assigned_agent,
            task_type=subtask.task_type,
            priority=original_task.priority,
            input_data=input_data,
            callback_url=None,
            created_at=original_task.created_at,
            timeout_seconds=int(self._subtask_timeout),
        )

        try:
            result = await asyncio.wait_for(
                agent.execute(sub_task_msg),
                timeout=self._subtask_timeout,
            )
            output = {
                "status": "completed",
                "output": result.output_data if hasattr(result, "output_data") else result,
            }

            # Publish progress via MessageBus if available
            if self._message_bus is not None:
                try:
                    from agentkit.bus.message import AgentMessage
                    await self._message_bus.publish(AgentMessage(
                        sender=subtask.assigned_agent,
                        recipient="orchestrator",
                        topic="task.progress",
                        payload={
                            "task_id": subtask.task_id,
                            "status": "completed",
                        },
                    ))
                except Exception as e:
                    logger.warning(f"Failed to publish progress via MessageBus: {e}")

            return output
        except asyncio.TimeoutError:
            error_result = {"status": "failed", "error": "Subtask timed out"}
            if self._message_bus is not None:
                try:
                    from agentkit.bus.message import AgentMessage
                    await self._message_bus.publish(AgentMessage(
                        sender=subtask.assigned_agent,
                        recipient="orchestrator",
                        topic="task.progress",
                        payload={
                            "task_id": subtask.task_id,
                            "status": "failed",
                            "error": "Subtask timed out",
                        },
                    ))
                except Exception as e:
                    logger.warning(f"Failed to publish progress via MessageBus: {e}")
            return error_result
        except Exception as e:
            error_result = {"status": "failed", "error": str(e)}
            if self._message_bus is not None:
                try:
                    from agentkit.bus.message import AgentMessage
                    await self._message_bus.publish(AgentMessage(
                        sender=subtask.assigned_agent,
                        recipient="orchestrator",
                        topic="task.progress",
                        payload={
                            "task_id": subtask.task_id,
                            "status": "failed",
                            "error": str(e),
                        },
                    ))
                except Exception as e:
                    logger.warning(f"Failed to publish progress via MessageBus: {e}")
            return error_result

    def _inject_dependency_results(
        self,
        subtask: SubTask,
        subtask_results: dict[str, dict[str, Any]],
    ) -> dict[str, Any]:
        """将依赖子任务的结果注入到当前子任务的输入中"""
        enriched = dict(subtask.input_data)

        if subtask.depends_on:
            dep_results = {}
            for dep_id in subtask.depends_on:
                if dep_id in subtask_results:
                    dep_results[dep_id] = subtask_results[dep_id]
            if dep_results:
                enriched["dependency_results"] = dep_results

        return enriched

    async def _aggregate_results(
        self,
        plan: OrchestrationPlan,
        subtask_results: dict[str, dict[str, Any]],
        original_task: TaskMessage,
    ) -> dict[str, Any]:
        """汇总子任务结果"""
        # Simple aggregation: collect all outputs
        outputs = {}
        errors = []

        for subtask in plan.subtasks:
            result = subtask_results.get(subtask.task_id, {})
            if result.get("status") == "completed":
                outputs[subtask.task_id] = result.get("output", {})
            else:
                errors.append({
                    "task_id": subtask.task_id,
                    "error": result.get("error", "Unknown error"),
                })

        aggregated = {
            "outputs": outputs,
            "task_id": original_task.task_id,
        }
        if errors:
            aggregated["errors"] = errors
            aggregated["partial_success"] = True

        return aggregated

    def _get_available_skill_names(self) -> list[str]:
        """获取可用 Skill 名称列表"""
        try:
            agents_info = self._agent_pool.list_agents()
            return [a["name"] for a in agents_info]
        except Exception:
            return []

    def _convert_execution_plan_to_subtasks(
        self,
        execution_plan: Any,
        parent_task_id: str,
        default_agent: str,
        default_task_type: str,
        original_input: dict[str, Any],
    ) -> list[SubTask]:
        """将 ExecutionPlan 的 PlanStep 转换为 SubTask 列表"""
        subtasks: list[SubTask] = []

        for step in execution_plan.steps:
            # 尝试根据 required_skills 匹配 agent
            assigned_agent = default_agent
            if step.required_skills:
                matched_agent = self._match_agent_for_skills(step.required_skills)
                if matched_agent:
                    assigned_agent = matched_agent

            subtasks.append(SubTask(
                task_id=step.step_id,
                parent_task_id=parent_task_id,
                assigned_agent=assigned_agent,
                task_type=default_task_type,
                input_data={
                    **original_input,
                    "step_name": step.name,
                    "step_description": step.description,
                },
                depends_on=list(step.dependencies),
            ))

        return subtasks

    def _match_agent_for_skills(self, required_skills: list[str]) -> str | None:
        """根据所需 Skill 匹配 Agent"""
        try:
            agents_info = self._agent_pool.list_agents()
            for skill in required_skills:
                for agent in agents_info:
                    name = agent.get("name", "")
                    agent_type = agent.get("agent_type", "")
                    description = agent.get("description", "").lower()
                    if skill.lower() in name.lower() or skill.lower() in agent_type.lower() or skill.lower() in description:
                        return name
        except Exception:
            pass
        return None

    async def execute_adaptive(
        self, task: TaskMessage,
    ) -> OrchestrationResult:
        """自适应编排：执行→评估→再分解循环。

        与 execute() 不同，此方法在第一轮执行后评估子任务结果质量，
        如果评估不通过且未达 max_iterations，则基于评估反馈重新分解
        未达标的子任务，保留已完成的子任务结果，然后执行新分解的子任务。

        Args:
            task: 原始任务消息

        Returns:
            OrchestrationResult: 编排结果，metadata 中包含迭代历史
        """
        import time as _time

        start_time = _time.monotonic()
        iteration_history: list[dict[str, Any]] = []

        # First execution
        result = await self.execute(task)

        # If adaptive not enabled or already succeeded, return directly
        if not self._config.adaptive or result.status == TaskStatus.COMPLETED:
            # Check quality even on success
            if self._config.adaptive and self._llm_gateway:
                quality = await self._evaluate_quality(task, result)
                if quality["score"] >= self._config.quality_threshold:
                    result.metadata["quality_score"] = quality["score"]
                    return result
            return result

        # Adaptive loop
        current_result = result
        for iteration in range(1, self._config.max_iterations + 1):
            # Evaluate quality
            quality = await self._evaluate_quality(task, current_result)
            iteration_history.append({
                "iteration": iteration,
                "quality_score": quality["score"],
                "feedback": quality.get("feedback", ""),
            })

            if quality["score"] >= self._config.quality_threshold:
                logger.info(
                    f"Adaptive iteration {iteration}: quality "
                    f"{quality['score']:.2f} >= {self._config.quality_threshold}"
                )
                current_result.metadata["quality_score"] = quality["score"]
                current_result.metadata["iterations"] = iteration_history
                return current_result

            logger.info(
                f"Adaptive iteration {iteration}: quality "
                f"{quality['score']:.2f} < {self._config.quality_threshold}, "
                f"re-decomposing failed subtasks"
            )

            # Re-decompose failed subtasks
            new_result = await self._reexecute_failed(
                task, current_result, quality,
            )
            current_result = new_result

        # Exhausted iterations
        current_result.metadata["iterations"] = iteration_history
        return current_result

    async def _evaluate_quality(
        self,
        task: TaskMessage,
        result: OrchestrationResult,
    ) -> dict[str, Any]:
        """评估子任务结果质量。

        Returns:
            Dict with "score" (0-1) and optional "feedback" string.
        """
        # Rule-based evaluation when no LLM
        if self._llm_gateway is None:
            return self._rule_based_evaluate(result)

        try:
            return await self._llm_evaluate(task, result)
        except Exception as e:
            logger.warning(f"LLM evaluation failed, falling back to rule-based: {e}")
            return self._rule_based_evaluate(result)

    def _rule_based_evaluate(
        self, result: OrchestrationResult,
    ) -> dict[str, Any]:
        """基于规则的质量评估：根据完成率打分。"""
        total = len(result.subtask_results)
        if total == 0:
            return {"score": 0.0, "feedback": "No subtasks executed"}

        completed = sum(
            1 for r in result.subtask_results.values()
            if r.get("status") == "completed"
        )
        score = completed / total
        feedback = ""
        if score < 1.0:
            failed = [
                tid for tid, r in result.subtask_results.items()
                if r.get("status") != "completed"
            ]
            feedback = f"Failed subtasks: {failed}"
        return {"score": score, "feedback": feedback}

    async def _llm_evaluate(
        self,
        task: TaskMessage,
        result: OrchestrationResult,
    ) -> dict[str, Any]:
        """使用 LLM 评估子任务结果质量。"""
        import json

        subtask_summary = []
        for tid, r in result.subtask_results.items():
            subtask_summary.append({
                "task_id": tid,
                "status": r.get("status", "unknown"),
                "output_preview": str(r.get("output", ""))[:200],
            })

        prompt = (
            f"Evaluate the quality of the following orchestration result.\n\n"
            f"Original task: {task.input_data}\n"
            f"Subtask results:\n{json.dumps(subtask_summary, ensure_ascii=False)}\n\n"
            f'Respond ONLY with JSON: {{"score": 0.0-1.0, "feedback": "..."}}\n'
            f"Score 1.0 = perfect, 0.0 = completely failed."
        )

        response = await self._llm_gateway.chat(
            messages=[{"role": "user", "content": prompt}],
            model="default",
        )

        try:
            text = response.content.strip()
            if text.startswith("```"):
                lines = text.split("\n")
                text = "\n".join(lines[1:-1])
            data = json.loads(text)
            return {
                "score": float(data.get("score", 0.0)),
                "feedback": data.get("feedback", ""),
            }
        except (json.JSONDecodeError, ValueError) as e:
            logger.warning(f"Failed to parse LLM evaluation: {e}")
            return self._rule_based_evaluate(result)

    async def _reexecute_failed(
        self,
        task: TaskMessage,
        previous_result: OrchestrationResult,
        quality: dict[str, Any],
    ) -> OrchestrationResult:
        """重新执行失败的子任务，保留已完成的结果。"""
        import time as _time

        start_time = _time.monotonic()

        # Identify failed subtasks
        failed_task_ids = [
            tid for tid, r in previous_result.subtask_results.items()
            if r.get("status") != "completed"
        ]

        if not failed_task_ids:
            return previous_result

        # Create new subtasks for failed ones, incorporating feedback
        new_subtasks = []
        for tid in failed_task_ids:
            old_result = previous_result.subtask_results[tid]
            new_subtasks.append(SubTask(
                task_id=f"retry-{tid}",
                parent_task_id=task.task_id,
                assigned_agent=task.agent_name,
                task_type=task.task_type,
                input_data={
                    **task.input_data,
                    "previous_error": old_result.get("error", ""),
                    "improvement_feedback": quality.get("feedback", ""),
                },
            ))

        # Build a mini-plan for the retry subtasks
        plan = OrchestrationPlan(
            plan_id=f"retry-{previous_result.plan_id}",
            parent_task_id=task.task_id,
            subtasks=new_subtasks,
            parallel_groups=[[st.task_id for st in new_subtasks]],
        )

        # Execute retry subtasks
        retry_results = await self._execute_plan(plan, task)

        # Merge: keep completed results, replace failed with retry results
        merged_results = {}
        for tid, r in previous_result.subtask_results.items():
            if r.get("status") == "completed":
                merged_results[tid] = r

        for tid, r in retry_results.items():
            # Map retry task IDs back to original
            original_tid = tid.replace("retry-", "", 1)
            merged_results[original_tid] = r

        # Re-aggregate
        all_subtasks = []
        for tid, r in merged_results.items():
            all_subtasks.append(SubTask(
                task_id=tid,
                parent_task_id=task.task_id,
                assigned_agent=task.agent_name,
                task_type=task.task_type,
                input_data=task.input_data,
                status=SubTaskStatus.COMPLETED if r.get("status") == "completed" else SubTaskStatus.FAILED,
                result=r.get("output"),
            ))

        retry_plan = OrchestrationPlan(
            plan_id=plan.plan_id,
            parent_task_id=task.task_id,
            subtasks=all_subtasks,
            parallel_groups=[],
        )

        aggregated = await self._aggregate_results(retry_plan, merged_results, task)

        failed_count = sum(
            1 for r in merged_results.values() if r.get("status") == "failed"
        )
        if failed_count == len(merged_results):
            status = TaskStatus.FAILED
        elif failed_count > 0:
            status = TaskStatus.PARTIALLY_COMPLETED
        else:
            status = TaskStatus.COMPLETED

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

        return OrchestrationResult(
            plan_id=plan.plan_id,
            parent_task_id=task.task_id,
            subtask_results=merged_results,
            aggregated_result=aggregated,
            status=status,
            total_duration_ms=duration_ms,
        )
