"""AgentRearrange graph builder using arrow/comma flow syntax."""

from __future__ import annotations

from typing import Any, Iterable

from beaver.coordinator.models import ExecutionGraph

from .base import (
    WorkflowAgentSpec,
    agent_name_set,
    build_graph_from_dependencies,
    edges_to_dependencies,
    parse_agents,
    validate_no_disconnected_agents,
)

WORKFLOW_NAME = "AgentRearrange"


def build_graph(
    *,
    task: str,
    agents: Iterable[WorkflowAgentSpec | dict[str, Any]],
    flow: str,
) -> ExecutionGraph:
    del task
    parsed = parse_agents(agents)
    edges = parse_flow(flow, known_agents=agent_name_set(parsed))
    dependencies = edges_to_dependencies(agents=parsed, edges=edges)
    validate_no_disconnected_agents(agents=parsed, dependencies=dependencies)
    return build_graph_from_dependencies(
        workflow_name=WORKFLOW_NAME,
        strategy="dag",
        agents=parsed,
        dependencies=dependencies,
    )


def parse_flow(flow: str, *, known_agents: set[str]) -> list[tuple[str, str]]:
    stages = _parse_stages(flow)
    edges: list[tuple[str, str]] = []
    for stage in stages:
        for name in stage:
            if name not in known_agents:
                raise ValueError(f"workflow flow references unknown agent: {name}")
    for left, right in zip(stages, stages[1:], strict=False):
        for source in left:
            for target in right:
                edge = (source, target)
                if edge not in edges:
                    edges.append(edge)
    return edges


def _parse_stages(flow: str) -> list[list[str]]:
    raw_flow = str(flow or "").strip()
    if not raw_flow:
        raise ValueError("workflow flow is required")
    stages: list[list[str]] = []
    for raw_stage in raw_flow.split("->"):
        names = [name.strip() for name in raw_stage.split(",") if name.strip()]
        if not names:
            raise ValueError("workflow flow contains an empty stage")
        if len(names) != len(set(names)):
            raise ValueError("workflow flow contains duplicate agent names in a stage")
        stages.append(names)
    if len(stages) < 2:
        raise ValueError("workflow flow must contain at least two stages")
    return stages