"""Core models for Beaver team coordination."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal

if TYPE_CHECKING:
    from beaver.engine.context import SkillContext
    from beaver.tasks.evidence import RunEvidence


TeamStrategy = Literal[
    "sequence",
    "parallel",
    "dag",
    "moa",
    "hierarchy",
    "heavy",
    "group_chat",
    "forest",
    "maker",
    "router",
]


@dataclass(slots=True)
class AgentDescriptor:
    """Runtime identity for a delegated local agent."""

    name: str
    role: str = ""
    system_prompt: str = ""
    model: str | None = None
    provider_name: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class DelegationEnvelope:
    """All context passed from a parent agent run to one delegated run."""

    parent_task_id: str | None
    parent_session_id: str
    parent_run_id: str | None
    agent: AgentDescriptor
    task: str
    inherited_pinned_skills: list[str] = field(default_factory=list)
    inherited_pinned_skill_contexts: list["SkillContext"] = field(default_factory=list)
    constraints: list[str] = field(default_factory=list)
    expected_output: str | None = None
    node_id: str | None = None
    dependency_outputs: dict[str, str] = field(default_factory=dict)


@dataclass(slots=True)
class ExecutionNode:
    """One node in a team execution graph."""

    node_id: str
    task: str
    agent: AgentDescriptor
    depends_on: list[str] = field(default_factory=list)
    inherited_pinned_skills: list[str] = field(default_factory=list)
    inherited_pinned_skill_contexts: list["SkillContext"] = field(default_factory=list)
    constraints: list[str] = field(default_factory=list)
    expected_output: str | None = None


@dataclass(slots=True)
class ExecutionGraph:
    """A lightweight team graph built from Beaver-native execution nodes."""

    strategy: TeamStrategy
    nodes: list[ExecutionNode]

    def validate(self) -> None:
        if self.strategy not in {"sequence", "parallel", "dag"}:
            raise NotImplementedError(f"Team strategy {self.strategy!r} is reserved but not implemented in v1")
        if not self.nodes:
            raise ValueError("ExecutionGraph requires at least one node")
        node_ids = [node.node_id for node in self.nodes]
        if len(node_ids) != len(set(node_ids)):
            raise ValueError("ExecutionGraph node_id values must be unique")
        known = set(node_ids)
        for node in self.nodes:
            missing = [item for item in node.depends_on if item not in known]
            if missing:
                raise ValueError(f"ExecutionNode {node.node_id!r} depends on unknown node(s): {missing}")
        visiting: set[str] = set()
        visited: set[str] = set()
        deps = {node.node_id: list(node.depends_on) for node in self.nodes}

        def visit(node_id: str) -> None:
            if node_id in visited:
                return
            if node_id in visiting:
                raise ValueError(f"ExecutionGraph has cyclic or unresolved dependencies involving {node_id!r}")
            visiting.add(node_id)
            for dep in deps[node_id]:
                visit(dep)
            visiting.remove(node_id)
            visited.add(node_id)

        for node_id in node_ids:
            visit(node_id)


@dataclass(slots=True)
class NodeRunResult:
    """Normalized result for one team node."""

    node_id: str
    success: bool
    output_text: str
    run_id: str | None = None
    session_id: str | None = None
    finish_reason: str = "stop"
    error: str | None = None
    evidence: "RunEvidence | None" = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "node_id": self.node_id,
            "success": self.success,
            "output_text": self.output_text,
            "run_id": self.run_id,
            "session_id": self.session_id,
            "finish_reason": self.finish_reason,
            "error": self.error,
            "evidence": self.evidence.to_dict() if self.evidence is not None else None,
        }


@dataclass(slots=True)
class TeamRunResult:
    """Normalized result returned by a Beaver team run."""

    success: bool
    summary: str
    node_results: list[NodeRunResult] = field(default_factory=list)
    run_ids: list[str] = field(default_factory=list)
    session_ids: list[str] = field(default_factory=list)
    task_id: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "success": self.success,
            "summary": self.summary,
            "node_results": [item.to_dict() for item in self.node_results],
            "run_ids": list(self.run_ids),
            "session_ids": list(self.session_ids),
            "task_id": self.task_id,
        }
