from __future__ import annotations

"""Directed graph support for multi-agent routing."""

from dataclasses import dataclass, field

from agentrspy.core import AgentError, AgentOutput


@dataclass(slots=True)
class EdgeCondition:
    """Transition predicate for graph orchestration."""

    kind: str
    keyword: str | None = None

    @classmethod
    def always(cls) -> "EdgeCondition":
        return cls("always")

    @classmethod
    def contains(cls, keyword: str) -> "EdgeCondition":
        return cls("contains", keyword=keyword)

    @classmethod
    def end(cls) -> "EdgeCondition":
        return cls("end")


EdgeCondition.Always = EdgeCondition.always()  # type: ignore[attr-defined]
EdgeCondition.Contains = EdgeCondition.contains  # type: ignore[attr-defined]
EdgeCondition.End = EdgeCondition.end()  # type: ignore[attr-defined]


@dataclass(slots=True)
class GraphEdge:
    """Graph edge between two agent nodes."""

    from_node: str
    to: str
    condition: EdgeCondition


@dataclass(slots=True)
class AgentGraph:
    """Directed agent graph."""

    nodes: dict[str, str] = field(default_factory=dict)
    edges: list[GraphEdge] = field(default_factory=list)
    entry_node: str | None = None

    @classmethod
    def builder(cls) -> "AgentGraphBuilder":
        return AgentGraphBuilder()

    def entry(self) -> str:
        if self.entry_node is None:
            raise AgentError.invalid_configuration("graph entry point not set")
        return self.entry_node

    def next(self, current: str, output: AgentOutput) -> str | None:
        for edge in self.edges:
            if edge.from_node != current:
                continue
            if edge.condition.kind == "always":
                return edge.to
            if (
                edge.condition.kind == "contains"
                and edge.condition.keyword
                and edge.condition.keyword in output.text
            ):
                return edge.to
            if edge.condition.kind == "end":
                continue
        return None


@dataclass(slots=True)
class AgentGraphBuilder:
    """Builder for `AgentGraph`."""

    graph: AgentGraph = field(default_factory=AgentGraph)

    def node(self, name: str) -> "AgentGraphBuilder":
        self.graph.nodes[name] = name
        return self

    def edge(
        self, from_node: str, to: str, condition: EdgeCondition
    ) -> "AgentGraphBuilder":
        self.graph.edges.append(
            GraphEdge(from_node=from_node, to=to, condition=condition)
        )
        return self

    def entry(self, entry: str) -> "AgentGraphBuilder":
        self.graph.entry_node = entry
        return self

    def build(self) -> AgentGraph:
        entry = self.graph.entry()
        if entry not in self.graph.nodes:
            raise AgentError.invalid_configuration(
                f"graph entry node '{entry}' is not registered"
            )
        return self.graph
