"""Butler Task Orchestrator — DAG-based multi-agent task execution.

Uses Butler's own AgentLoop for sub-agent spawning. Every sub-agent
goes through Butler's orchestrator, inheriting the correct model
config, memory, skills, and agent profile.

Key improvements over v3:
  - AgentSpawnConfig includes tools and model_config fields
  - AgentResult includes AgentReport and tokens_used
  - Real parallel execution via asyncio.to_thread
  - Sub-agents use Butler profiles and toolsets
"""

from __future__ import annotations

import asyncio
from contextvars import copy_context
import logging
import threading
import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, cast

from butler.configuration.settings import ModelConfig
from butler.core.delegate_context import child_callbacks, get_parent_callbacks
from butler.core.meta_flags import MAX_DAG_NODES
from butler.dag_scheduler import (
    apply_node_router,
    finalize_unexecuted_nodes,
    graph_all_required_ok,
    group_into_layers,
    prepare_layer_node,
    split_layer_batches,
    topological_sort,
)
from butler.delegate.policy import DELEGATE_BLOCKED_TOOLS, MAX_DELEGATE_DEPTH
from butler.report import AgentReport, cache_report, enrich_report_decisions
from butler.session.lifecycle import attach_turn_memory_prefetch, sync_turn_memory
from butler.task_orchestrator_ops import (
    on_progress_safe,
    spawn_agent_loud,
    truncate_child_response_safe,
    workflow_max_parallel_safe,
)
from butler.tools.project_tools import canonical_tool_name, get_tool_definitions_for_project
from butler.tools.registry import _extract_changes_from_messages, _safe_dispatch
from butler.execution_context import (
    get_current_orchestrator,
    get_current_session_key,
    use_execution_context,
    use_workflow_step,
)

logger = logging.getLogger(__name__)
_MODEL_OVERRIDE_LOCK = threading.RLock()


class AgentStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class AgentSpawnConfig:
    """Configuration for spawning a sub-agent."""
    role: str  # dev, content, review
    task: str
    context: str = ""
    tools: list[str] | None = None
    model_config: dict[str, Any] | None = None
    max_iterations: int = 30
    output_format: str = "text"
    project_name: str | None = None
    delegate_depth: int = 0
    session_key: str = ""
    clear_child_transcript: bool = False


@dataclass
class AgentResult:
    """Result from a sub-agent execution, including structured report."""
    success: bool
    response: str = ""
    report: AgentReport | None = None
    tokens_used: int = 0
    iterations: int = 0
    tool_calls: int = 0
    elapsed_seconds: float = 0.0
    error: str = ""
    artifacts: list[str] = field(default_factory=list)


@dataclass
class AgentTask:
    id: str
    config: AgentSpawnConfig
    status: AgentStatus = AgentStatus.PENDING
    result: AgentResult | None = None
    started_at: float = 0.0
    completed_at: float = 0.0


@dataclass
class TaskNode:
    """Node in a task DAG."""
    id: str
    config: AgentSpawnConfig
    depends_on: list[str] = field(default_factory=list)
    # Return a direct dependent id to choose one branch; return None to fan out all dependents.
    router: Callable[[AgentResult], str | None] | None = None
    requires_approval: bool = False
    max_retries: int = 1
    handoff_only: bool = True
    clear_child_transcript: bool = False
    supervisor_note: str = ""
    optional: bool = False
    rescue_configs: list[AgentSpawnConfig] = field(default_factory=list)
    until: dict[str, Any] | None = None


@dataclass
class TaskGraphResult:
    nodes: dict[str, AgentResult] = field(default_factory=dict)
    execution_order: list[str] = field(default_factory=list)
    success: bool = True
    error: str = ""


class TaskOrchestrator:
    """Execute multi-agent workflows using Butler's AgentLoop."""

    def __init__(self) -> None:
        import threading

        self._tasks_lock = threading.Lock()
        self._tasks: dict[str, AgentTask] = {}

    def _create_agent_loop(self, config: AgentSpawnConfig) -> Any:
        """Create a Butler AgentLoop for the given config."""
        orch = _orchestrator_for_task()

        if config.model_config:
            mc = ModelConfig(**config.model_config)
            settings = orch._settings
            with _MODEL_OVERRIDE_LOCK:
                previous = getattr(settings, "_runtime_model_overrides", {}).get(config.role)
                settings.set_runtime_model_override(config.role, mc)
                try:
                    return self._create_agent_loop_with_orchestrator(config, orch)
                finally:
                    settings.set_runtime_model_override(config.role, previous)

        return self._create_agent_loop_with_orchestrator(config, orch)

    def _create_agent_loop_with_orchestrator(self, config: AgentSpawnConfig, orch: Any) -> Any:
        """Create a Butler AgentLoop using an already-selected orchestrator."""
        project = orch.project_manager.get_current(session_key=str(config.session_key or ""))
        tools = get_tool_definitions_for_project(project, role=config.role)
        delegated_tools = [
            t for t in tools
            if t["function"]["name"] not in DELEGATE_BLOCKED_TOOLS
        ]

        if config.tools:
            tool_names = {canonical_tool_name(n) for n in config.tools}
            delegated_tools = [
                t for t in delegated_tools
                if t["function"]["name"] in tool_names
            ]

        agent = orch.create_project_agent_loop(
            role=config.role,
            tools=delegated_tools,
            tool_dispatcher=lambda name, args: dispatch_tool_safely(
                name,
                args,
                depth=config.delegate_depth + 1,
            ),
            callbacks=child_callbacks(get_parent_callbacks()),
        )

        agent.config.max_iterations = config.max_iterations
        agent._butler_orchestrator = orch
        return agent

    async def spawn_agent(
        self,
        config: AgentSpawnConfig,
        on_progress: Callable[[str, str, str], None] | None = None,
    ) -> AgentResult:
        """Spawn a Butler AgentLoop and run the task."""
        task_id = uuid.uuid4().hex[:8]
        task = AgentTask(id=task_id, config=config)
        with self._tasks_lock:
            if len(self._tasks) > 200:
                completed = [
                    k for k, v in self._tasks.items()
                    if getattr(v, "status", None) in ("completed", "error", "cancelled")
                ]
                for k in completed[:100]:
                    del self._tasks[k]
            self._tasks[task_id] = task
        task.status = AgentStatus.RUNNING
        task.started_at = time.time()

        logger.info("Spawning %s agent [%s]: %s", config.role, task_id, config.task[:80])

        on_progress_safe(on_progress, task_id, "start", config.role)

        async def _run_spawn() -> AgentResult:
            if config.delegate_depth >= MAX_DELEGATE_DEPTH:
                raise ValueError(f"Maximum delegation depth ({MAX_DELEGATE_DEPTH}) exceeded")

            agent = self._create_agent_loop(config)

            raw_user_message = config.task
            if config.context:
                raw_user_message = f"## 上下文\n{config.context}\n\n## 任务\n{config.task}"

            agent.reset()
            orch = getattr(agent, "_butler_orchestrator", None)
            session_key = _resolve_spawn_session_key(config, task_id)
            user_message = raw_user_message
            if orch is not None and hasattr(orch, "inject_skill_context"):
                user_message = orch.inject_skill_context(raw_user_message)
                attach_turn_memory_prefetch(agent, orch, raw_user_message, role=config.role)
            with use_execution_context(orch, session_key=session_key):
                run_context = copy_context()
                loop_result = await asyncio.to_thread(
                    run_context.run,
                    agent.run,
                    user_message,
                )
            if orch is not None:
                sync_turn_memory(
                    orch,
                    raw_user_message,
                    loop_result.final_response or "",
                    interrupted=loop_result.status.value == "interrupted",
                    status=loop_result.status,
                    session_id=session_key,
                )

            changes = _extract_changes_from_messages(loop_result.messages)
            report = AgentReport(
                headline=f"{config.role} 代理完成任务",
                summary=loop_result.final_response or "",
                changes=changes,
                success=loop_result.status.value == "completed",
                iterations=loop_result.iterations,
                tool_calls=loop_result.tool_calls_made,
                tokens_used=loop_result.total_tokens,
                elapsed_seconds=loop_result.elapsed_seconds,
            )
            enrich_report_decisions(report, loop_result.final_response or "")
            cache_report(report, session_key=session_key)
            success = loop_result.status.value == "completed"

            response_text = loop_result.final_response or ""
            response_text = truncate_child_response_safe(
                response_text=response_text,
                report=report,
                clear_child_transcript=config.clear_child_transcript,
            )
            return AgentResult(
                success=success,
                response=response_text,
                report=report,
                tokens_used=loop_result.total_tokens,
                iterations=loop_result.iterations,
                tool_calls=loop_result.tool_calls_made,
                elapsed_seconds=loop_result.elapsed_seconds,
                error="" if success else (loop_result.error or loop_result.status.value),
            )

        result = await spawn_agent_loud(_run_spawn, task_id=task_id)

        task.completed_at = time.time()
        task.status = AgentStatus.COMPLETED if result.success else AgentStatus.FAILED
        task.result = result

        elapsed = task.completed_at - task.started_at
        logger.info("Agent [%s] %s in %.1fs", task_id, task.status.value, elapsed)

        if on_progress:
            status_label = "done" if result.success else "failed"
            preview = ""
            if status_label != "start":
                preview = (result.response or result.error or "")[:500]
            on_progress_safe(
                on_progress,
                task_id,
                status_label,
                config.role,
                preview,
            )

        return cast(AgentResult, result)

    async def spawn_parallel(self, configs: list[AgentSpawnConfig]) -> list[AgentResult]:
        """Spawn multiple agents in true parallel via asyncio.to_thread."""
        tasks = [self.spawn_agent(cfg) for cfg in configs]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        return [_coerce_agent_result(result) for result in results]

    async def spawn_sequential(
        self,
        configs: list[AgentSpawnConfig],
        pass_context: bool = True,
    ) -> list[AgentResult]:
        """Execute agents sequentially, optionally passing context."""
        results: list[AgentResult] = []
        accumulated_context = ""

        for cfg in configs:
            if pass_context and accumulated_context:
                cfg.context = (cfg.context or "") + "\n\n" + accumulated_context

            result = await self.spawn_agent(cfg)
            results.append(result)

            if result.success and result.response:
                accumulated_context += f"\n[{cfg.role} 结果]: {result.response[:2000]}"

            if not result.success:
                logger.warning("Sequential chain broken at role=%s", cfg.role)
                break

        return results

    async def execute_graph(
        self,
        nodes: list[TaskNode],
        on_approval: Callable[[TaskNode], bool] | None = None,
        on_progress: Callable[[str, str, str], None] | None = None,
        *,
        max_parallel: int | None = None,
        serial: bool = False,
    ) -> TaskGraphResult:
        """Execute a DAG of tasks with dependency resolution."""
        if len(nodes) > MAX_DAG_NODES:
            raise ValueError(
                f"DAG exceeds maximum node count ({len(nodes)} > {MAX_DAG_NODES})"
            )

        graph_result = TaskGraphResult()
        node_map = {n.id: n for n in nodes}
        completed: dict[str, AgentResult] = {}
        cancelled: set[str] = set()
        errors: list[str] = []

        order = topological_sort(nodes)

        layers = group_into_layers(order, node_map)
        step_index = 0

        if max_parallel is None:
            max_parallel = workflow_max_parallel_safe(default=None)
        for layer in layers:
            layer_tasks = []
            for node_id in layer:
                node = node_map[node_id]
                if node_id in completed or node_id in cancelled:
                    continue

                skip = prepare_layer_node(
                    node,
                    completed=completed,
                    cancelled=cancelled,
                    node_map=node_map,
                    on_approval=on_approval,
                )
                if skip is not None:
                    completed[node_id] = skip
                    continue

                layer_tasks.append((node_id, node))

            if not layer_tasks:
                continue

            batches = split_layer_batches(
                layer_tasks, serial=serial, max_parallel=max_parallel,
            )

            for batch in batches:
                if len(batch) == 1:
                    nid, node = batch[0]
                    step_index += 1
                    on_progress_safe(on_progress, nid, "start", node.config.role)
                    with use_workflow_step(nid):
                        result = _coerce_agent_result(
                            await self._run_with_retry(node, on_progress=on_progress)
                        )
                    completed[nid] = result
                    graph_result.execution_order.append(nid)
                else:
                    async def _run_one(n: TaskNode) -> AgentResult:
                        with use_workflow_step(n.id):
                            return await self._run_with_retry(n, on_progress=on_progress)

                    parallel_results = await asyncio.gather(
                        *[_run_one(node) for _, node in batch],
                        return_exceptions=True,
                    )
                    for (nid, _), raw_parallel in zip(batch, parallel_results):
                        completed[nid] = _coerce_agent_result(raw_parallel)
                        graph_result.execution_order.append(nid)

            for nid in [lid for lid, _ in layer_tasks]:
                node = node_map[nid]
                result = completed[nid]
                errors.extend(
                    apply_node_router(
                        nid, node, result, node_map=node_map, cancelled=cancelled,
                    )
                )

        finalize_unexecuted_nodes(
            order,
            completed=completed,
            cancelled=cancelled,
            node_map=node_map,
            errors=errors,
        )

        graph_result.nodes = completed
        graph_result.error = "; ".join(errors)
        graph_result.success = (
            not errors and graph_all_required_ok(completed, node_map)
        )
        return graph_result

    async def _run_with_retry(
        self,
        node: TaskNode,
        on_progress: Callable[[str, str, str], None] | None = None,
    ) -> AgentResult:
        from butler.workflows.step_runner import run_step_with_retry

        return cast(
            AgentResult,
            await run_step_with_retry(
                node,
                spawn=self.spawn_agent,
                on_progress=on_progress,
            ),
        )

    async def _run_rescue_steps(
        self,
        node: TaskNode,
        failed: AgentResult,
        *,
        on_progress: Callable[[str, str, str], None] | None = None,
    ) -> AgentResult:
        from butler.workflows.step_runner import run_rescue_steps

        return cast(
            AgentResult,
            await run_rescue_steps(
                node,
                failed,
                spawn=self.spawn_agent,
                on_progress=on_progress,
            ),
        )


def _resolve_spawn_session_key(config: AgentSpawnConfig, task_id: str) -> str:
    """Pick a stable audit/session key for orchestrator spawns."""
    explicit = str(config.session_key or "").strip()
    if explicit:
        return explicit
    inherited = str(get_current_session_key() or "").strip()
    if inherited:
        return inherited
    return f"task:{task_id}"


def _orchestrator_for_task() -> Any:
    orch = get_current_orchestrator()
    if orch is not None:
        return orch

    from butler.orchestrator import ButlerOrchestrator

    return ButlerOrchestrator(user_id="owner", channel="orchestrator")


def dispatch_tool_safely(name: str, args: dict[str, Any], *, depth: int) -> str:
    """Dispatch a tool call with delegate depth propagated consistently."""
    return str(_safe_dispatch(name, args, depth))


def _coerce_agent_result(result: AgentResult | BaseException) -> AgentResult:
    if isinstance(result, AgentResult):
        return result
    return AgentResult(success=False, error=str(result))


# Backward-compat aliases — canonical: ``butler.dag_scheduler`` (ENG-4)