"""Minimal scheduler for Beaver-native team execution graphs."""

from __future__ import annotations

import asyncio
from collections.abc import Callable
from typing import TYPE_CHECKING

from beaver.engine.providers import ProviderBundle

from ..local import LocalAgentRunner
from ..models import DelegationEnvelope, ExecutionGraph, ExecutionNode, NodeRunResult, TeamRunResult

if TYPE_CHECKING:
    from beaver.engine.context import SkillContext


class TeamGraphScheduler:
    """Execute sequence, parallel, and DAG team graphs."""

    def __init__(self, runner: LocalAgentRunner, *, max_parallel_team_nodes: int = 3) -> None:
        self.runner = runner
        self.max_parallel_team_nodes = max(1, int(max_parallel_team_nodes))

    async def run(
        self,
        graph: ExecutionGraph,
        *,
        parent_task_id: str | None,
        parent_session_id: str,
        parent_run_id: str | None = None,
        provider_bundle: ProviderBundle | None = None,
        provider_bundle_factory: Callable[[ExecutionNode], ProviderBundle | None] | None = None,
        inherited_pinned_skills: list[str] | None = None,
        inherited_pinned_skill_contexts: list["SkillContext"] | None = None,
        allow_candidate_generation: bool = False,
    ) -> TeamRunResult:
        graph.validate()
        if provider_bundle is not None and len(graph.nodes) > 1:
            raise ValueError("provider_bundle can only be used for single-node team graphs; use provider_bundle_factory")
        inherited = list(inherited_pinned_skills or [])
        inherited_contexts = list(inherited_pinned_skill_contexts or [])
        if graph.strategy == "sequence":
            results = await self._run_sequence(
                graph.nodes,
                parent_task_id=parent_task_id,
                parent_session_id=parent_session_id,
                parent_run_id=parent_run_id,
                provider_bundle=provider_bundle,
                provider_bundle_factory=provider_bundle_factory,
                inherited_pinned_skills=inherited,
                inherited_pinned_skill_contexts=inherited_contexts,
                allow_candidate_generation=allow_candidate_generation,
            )
        elif graph.strategy == "parallel":
            results = await self._run_parallel(
                graph.nodes,
                parent_task_id=parent_task_id,
                parent_session_id=parent_session_id,
                parent_run_id=parent_run_id,
                provider_bundle=provider_bundle,
                provider_bundle_factory=provider_bundle_factory,
                inherited_pinned_skills=inherited,
                inherited_pinned_skill_contexts=inherited_contexts,
                allow_candidate_generation=allow_candidate_generation,
            )
        else:
            results = await self._run_dag(
                graph.nodes,
                parent_task_id=parent_task_id,
                parent_session_id=parent_session_id,
                parent_run_id=parent_run_id,
                provider_bundle=provider_bundle,
                provider_bundle_factory=provider_bundle_factory,
                inherited_pinned_skills=inherited,
                inherited_pinned_skill_contexts=inherited_contexts,
                allow_candidate_generation=allow_candidate_generation,
            )
        return self._summarize(results, task_id=parent_task_id)

    async def _run_sequence(
        self,
        nodes: list[ExecutionNode],
        **kwargs,
    ) -> list[NodeRunResult]:
        results: list[NodeRunResult] = []
        for node in nodes:
            if any(not item.success for item in results):
                results.append(self._blocked(node, results))
                continue
            dependency_outputs = {item.node_id: item.output_text for item in results if item.success}
            results.append(await self._run_node(node, dependency_outputs=dependency_outputs, **kwargs))
        return results

    async def _run_parallel(
        self,
        nodes: list[ExecutionNode],
        **kwargs,
    ) -> list[NodeRunResult]:
        semaphore = asyncio.Semaphore(self.max_parallel_team_nodes)

        async def run_one(node: ExecutionNode) -> NodeRunResult:
            async with semaphore:
                return await self._run_node(
                    node,
                    dependency_outputs={},
                    execution_mode="isolated_loop",
                    **kwargs,
                )

        return list(await asyncio.gather(*(run_one(node) for node in nodes)))

    async def _run_dag(
        self,
        nodes: list[ExecutionNode],
        **kwargs,
    ) -> list[NodeRunResult]:
        pending = {node.node_id: node for node in nodes}
        completed: dict[str, NodeRunResult] = {}
        ordered: list[NodeRunResult] = []

        while pending:
            blocked_ids = {
                node_id
                for node_id, node in pending.items()
                if any(dep in completed and not completed[dep].success for dep in node.depends_on)
            }
            for node_id in sorted(blocked_ids):
                node = pending.pop(node_id)
                result = self._blocked(node, list(completed.values()))
                completed[node_id] = result
                ordered.append(result)

            ready = [
                node
                for node in pending.values()
                if all(dep in completed and completed[dep].success for dep in node.depends_on)
            ]
            if not ready:
                if pending:
                    unresolved = ", ".join(sorted(pending))
                    raise ValueError(f"ExecutionGraph has cyclic or unresolved dependencies: {unresolved}")
                break

            batch = await asyncio.gather(
                *(
                    self._run_node(
                        node,
                        dependency_outputs={
                            dep: completed[dep].output_text
                            for dep in node.depends_on
                            if dep in completed
                        },
                        **kwargs,
                    )
                    for node in ready
                )
            )
            for result in batch:
                pending.pop(result.node_id, None)
                completed[result.node_id] = result
                ordered.append(result)

        return ordered

    async def _run_node(
        self,
        node: ExecutionNode,
        *,
        parent_task_id: str | None,
        parent_session_id: str,
        parent_run_id: str | None,
        provider_bundle: ProviderBundle | None,
        provider_bundle_factory: Callable[[ExecutionNode], ProviderBundle | None] | None,
        inherited_pinned_skills: list[str],
        inherited_pinned_skill_contexts: list["SkillContext"],
        allow_candidate_generation: bool,
        dependency_outputs: dict[str, str],
        execution_mode: str = "shared_loop",
    ) -> NodeRunResult:
        try:
            pinned = self._merge_pinned(inherited_pinned_skills, node.inherited_pinned_skills)
            pinned_contexts = self._merge_skill_contexts(
                inherited_pinned_skill_contexts,
                node.inherited_pinned_skill_contexts,
            )
            envelope = DelegationEnvelope(
                parent_task_id=parent_task_id,
                parent_session_id=parent_session_id,
                parent_run_id=parent_run_id,
                agent=node.agent,
                task=node.task,
                inherited_pinned_skills=pinned,
                inherited_pinned_skill_contexts=pinned_contexts,
                constraints=list(node.constraints),
                expected_output=node.expected_output,
                node_id=node.node_id,
                dependency_outputs=dict(dependency_outputs),
            )
            node_provider_bundle = provider_bundle_factory(node) if provider_bundle_factory is not None else provider_bundle
            return await self.runner.run(
                envelope,
                provider_bundle=node_provider_bundle,
                allow_candidate_generation=allow_candidate_generation,
                execution_mode=execution_mode,
            )
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            return NodeRunResult(
                node_id=node.node_id,
                success=False,
                output_text="",
                finish_reason="error",
                error=str(exc),
            )

    @staticmethod
    def _merge_pinned(parent: list[str], local: list[str]) -> list[str]:
        result: list[str] = []
        for name in [*parent, *local]:
            if name and name not in result:
                result.append(name)
        return result

    @staticmethod
    def _merge_skill_contexts(parent: list["SkillContext"], local: list["SkillContext"]) -> list["SkillContext"]:
        result: list["SkillContext"] = []
        seen: set[str] = set()
        for skill in [*parent, *local]:
            name = getattr(skill, "name", "")
            if not name or name in seen:
                continue
            seen.add(name)
            result.append(skill)
        return result

    @staticmethod
    def _blocked(node: ExecutionNode, prior_results: list[NodeRunResult]) -> NodeRunResult:
        failed = [item.node_id for item in prior_results if not item.success]
        detail = ", ".join(failed) or "unknown dependency"
        return NodeRunResult(
            node_id=node.node_id,
            success=False,
            output_text="",
            finish_reason="blocked",
            error=f"Blocked by failed dependency: {detail}",
        )

    @staticmethod
    def _summarize(results: list[NodeRunResult], *, task_id: str | None) -> TeamRunResult:
        success = all(item.success for item in results)
        successful_outputs = [item.output_text.strip() for item in results if item.success and item.output_text.strip()]
        summary_parts = list(successful_outputs)
        failed = [item for item in results if not item.success]
        if failed:
            failure_lines = [
                f"- {item.node_id}: {item.error or item.finish_reason} evidence={'yes' if item.evidence else 'no'}"
                for item in failed
            ]
            summary_parts.append("Failed nodes:\n" + "\n".join(failure_lines))
        summary = "\n\n".join(summary_parts)
        return TeamRunResult(
            success=success,
            summary=summary,
            node_results=results,
            run_ids=[item.run_id for item in results if item.run_id],
            session_ids=[item.session_id for item in results if item.session_id],
            task_id=task_id,
        )
