"""Shared builders for local team workflow graph construction."""

from __future__ import annotations

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

from beaver.coordinator.models import AgentDescriptor, ExecutionGraph, ExecutionNode


GraphStrategy = Literal["sequence", "parallel", "dag"]


@dataclass(slots=True)
class WorkflowAgentSpec:
    name: str
    instruction: str
    use_skill: str | None = None
    skill_query: str | None = None
    allowed_tool_names: list[str] | None = None
    required_evidence: list[str] = field(default_factory=list)
    evidence_contract: dict[str, Any] = field(default_factory=dict)
    validation_rules: list[str] = field(default_factory=list)
    required_for_completion: bool = True
    block_downstream_on_partial: bool = False
    max_tool_iterations: int | None = None
    constraints: list[str] = field(default_factory=list)
    expected_output: str | None = None
    input_contract: dict[str, Any] = field(default_factory=dict)
    output_contract: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class WorkflowBuildResult:
    graph: ExecutionGraph
    workflow_name: str


def parse_agents(raw_agents: Iterable[WorkflowAgentSpec | dict[str, Any]]) -> list[WorkflowAgentSpec]:
    agents: list[WorkflowAgentSpec] = []
    for index, raw in enumerate(raw_agents, start=1):
        if isinstance(raw, WorkflowAgentSpec):
            spec = raw
        elif isinstance(raw, dict):
            spec = _agent_from_dict(raw, index=index)
        else:
            raise ValueError("workflow agents must be objects")
        agents.append(spec)
    validate_agent_names(agents)
    return agents


def validate_agent_names(agents: list[WorkflowAgentSpec]) -> None:
    if not agents:
        raise ValueError("workflow requires at least one agent")
    seen: set[str] = set()
    for agent in agents:
        if not agent.name:
            raise ValueError("workflow agent name is required")
        if not agent.instruction:
            raise ValueError(f"workflow agent {agent.name!r} requires instruction")
        if agent.name in seen:
            raise ValueError(f"workflow agent names must be unique: {agent.name}")
        seen.add(agent.name)


def agent_name_set(agents: list[WorkflowAgentSpec]) -> set[str]:
    return {agent.name for agent in agents}


def build_graph_from_dependencies(
    *,
    workflow_name: str,
    strategy: GraphStrategy,
    agents: list[WorkflowAgentSpec],
    dependencies: dict[str, list[str]],
) -> ExecutionGraph:
    nodes = [
        build_node(
            workflow_name=workflow_name,
            agent=agent,
            depends_on=dependencies.get(agent.name, []),
        )
        for agent in agents
    ]
    graph = ExecutionGraph(strategy=strategy, nodes=nodes)
    graph.validate()
    return graph


def build_node(
    *,
    workflow_name: str,
    agent: WorkflowAgentSpec,
    depends_on: list[str],
) -> ExecutionNode:
    metadata = {
        "sub_agent_kind": "generic_skill_worker",
        "workflow_tool": workflow_name,
        "workflow_agent_name": agent.name,
    }
    if agent.use_skill:
        metadata["use_skill"] = agent.use_skill
    if agent.skill_query:
        metadata["skill_query"] = agent.skill_query
    return ExecutionNode(
        node_id=agent.name,
        task=agent.instruction,
        agent=AgentDescriptor(
            name=agent.name,
            role="",
            system_prompt="",
            metadata=metadata,
        ),
        depends_on=list(depends_on),
        constraints=list(agent.constraints),
        expected_output=agent.expected_output,
        input_contract=dict(agent.input_contract),
        output_contract=dict(agent.output_contract),
        allowed_tool_names=(
            None if agent.allowed_tool_names is None else list(agent.allowed_tool_names)
        ),
        required_evidence=list(agent.required_evidence),
        evidence_contract=dict(agent.evidence_contract),
        validation_rules=list(agent.validation_rules),
        required_for_completion=agent.required_for_completion,
        block_downstream_on_partial=agent.block_downstream_on_partial,
        max_tool_iterations=agent.max_tool_iterations,
    )


def edges_to_dependencies(
    *,
    agents: list[WorkflowAgentSpec],
    edges: Iterable[tuple[str, str] | list[str]],
) -> dict[str, list[str]]:
    known = agent_name_set(agents)
    dependencies = {agent.name: [] for agent in agents}
    for raw_edge in edges:
        source, target = _parse_edge(raw_edge)
        if source not in known:
            raise ValueError(f"workflow edge references unknown agent: {source}")
        if target not in known:
            raise ValueError(f"workflow edge references unknown agent: {target}")
        if source == target:
            raise ValueError(f"workflow edge creates a self-cycle: {source}")
        if source not in dependencies[target]:
            dependencies[target].append(source)
    return dependencies


def validate_output_agent(
    *,
    agents: list[WorkflowAgentSpec],
    dependencies: dict[str, list[str]],
    output_agent: str,
    allow_disconnected: bool = False,
) -> None:
    known = agent_name_set(agents)
    if output_agent not in known:
        raise ValueError(f"workflow output_agent references unknown agent: {output_agent}")

    upstream = _upstream_nodes(output_agent, dependencies)
    if not upstream:
        raise ValueError(f"workflow output_agent {output_agent!r} must be reachable from upstream agents")

    if allow_disconnected:
        return

    connected = set(upstream)
    connected.add(output_agent)
    disconnected = sorted(known - connected)
    if disconnected:
        raise ValueError(f"workflow has disconnected agent(s): {', '.join(disconnected)}")


def validate_no_disconnected_agents(
    *,
    agents: list[WorkflowAgentSpec],
    dependencies: dict[str, list[str]],
) -> None:
    known = agent_name_set(agents)
    connected: set[str] = set()
    for target, sources in dependencies.items():
        if sources:
            connected.add(target)
            connected.update(sources)
    disconnected = sorted(known - connected)
    if disconnected:
        raise ValueError(f"workflow has disconnected agent(s): {', '.join(disconnected)}")


def _agent_from_dict(raw: dict[str, Any], *, index: int) -> WorkflowAgentSpec:
    name = _required_str(raw.get("name"), f"agents[{index}].name")
    instruction = _required_str(raw.get("instruction"), f"agents[{index}].instruction")
    return WorkflowAgentSpec(
        name=name,
        instruction=instruction,
        use_skill=_optional_str(raw.get("use_skill")),
        skill_query=_optional_str(raw.get("skill_query")),
        allowed_tool_names=_optional_string_list(raw.get("allowed_tool_names")),
        required_evidence=_string_list(raw.get("required_evidence")),
        evidence_contract=_dict(raw.get("evidence_contract")),
        validation_rules=_string_list(raw.get("validation_rules")),
        required_for_completion=bool(raw.get("required_for_completion", True)),
        block_downstream_on_partial=bool(raw.get("block_downstream_on_partial", False)),
        max_tool_iterations=_optional_int(raw.get("max_tool_iterations")),
        constraints=_string_list(raw.get("constraints")),
        expected_output=_optional_str(raw.get("expected_output")),
        input_contract=_dict(raw.get("input_contract")),
        output_contract=_dict(raw.get("output_contract")),
    )


def _parse_edge(raw_edge: tuple[str, str] | list[str]) -> tuple[str, str]:
    if not isinstance(raw_edge, (list, tuple)) or len(raw_edge) != 2:
        raise ValueError("workflow edges must be [source, target] pairs")
    source = _required_str(raw_edge[0], "edge source")
    target = _required_str(raw_edge[1], "edge target")
    return source, target


def _upstream_nodes(node_id: str, dependencies: dict[str, list[str]]) -> set[str]:
    result: set[str] = set()

    def visit(current: str) -> None:
        for dependency in dependencies.get(current, []):
            if dependency in result:
                continue
            result.add(dependency)
            visit(dependency)

    visit(node_id)
    return result


def _required_str(value: Any, label: str) -> str:
    text = str(value or "").strip()
    if not text:
        raise ValueError(f"{label} is required")
    return text


def _optional_str(value: Any) -> str | None:
    text = str(value or "").strip()
    return text or None


def _string_list(value: Any) -> list[str]:
    if value is None:
        return []
    if not isinstance(value, list):
        raise ValueError("expected a list of strings")
    return [str(item).strip() for item in value if str(item).strip()]


def _optional_string_list(value: Any) -> list[str] | None:
    if value is None:
        return None
    return _string_list(value)


def _dict(value: Any) -> dict[str, Any]:
    return dict(value) if isinstance(value, dict) else {}


def _optional_int(value: Any) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError) as exc:
        raise ValueError("max_tool_iterations must be an integer") from exc