from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable

START = "__start__"
END = "__end__"


@dataclass
class _CompiledGraph:
    state_type: Any
    nodes: dict[str, Callable[[dict[str, Any]], dict[str, Any]]]
    edges: dict[str, str]
    conditionals: dict[str, tuple[Callable[[dict[str, Any]], str], dict[str, str]]]

    def invoke(self, initial_state: dict[str, Any]) -> dict[str, Any]:
        current = self.edges[START]
        state = dict(initial_state)
        while current != END:
            state = self.nodes[current](state)
            if current in self.conditionals:
                router, mapping = self.conditionals[current]
                current = mapping[router(state)]
            else:
                current = self.edges[current]
        return state


class StateGraph:
    def __init__(self, state_type: Any):
        self.state_type = state_type
        self.nodes: dict[str, Callable[[dict[str, Any]], dict[str, Any]]] = {}
        self.edges: dict[str, str] = {}
        self.conditionals: dict[str, tuple[Callable[[dict[str, Any]], str], dict[str, str]]] = {}

    def add_node(self, name: str, fn: Callable[[dict[str, Any]], dict[str, Any]]) -> None:
        self.nodes[name] = fn

    def add_edge(self, source: str, dest: str) -> None:
        self.edges[source] = dest

    def add_conditional_edges(
        self,
        source: str,
        router: Callable[[dict[str, Any]], str],
        mapping: dict[str, str],
    ) -> None:
        self.conditionals[source] = (router, mapping)

    def compile(self) -> _CompiledGraph:
        return _CompiledGraph(
            state_type=self.state_type,
            nodes=self.nodes,
            edges=self.edges,
            conditionals=self.conditionals,
        )
