"""Benchmarks for PureGraph execution throughput and ordering."""

from __future__ import annotations

from typing import Any, ClassVar

from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
from cleveragents.langgraph.pure_graph import PureGraph


class PureGraphBench:
    """Benchmark suite for PureGraph execution performance."""

    params: ClassVar[list[int]] = [10, 50, 100, 500]
    param_names: ClassVar[list[str]] = ["node_count"]

    def setup(self, node_count: int) -> None:
        """Set up benchmark with varying node counts."""
        self.node_count = node_count
        self.nodes: dict[str, NodeConfig] = {
            "start": NodeConfig(name="start", type=NodeType.START),
            "end": NodeConfig(name="end", type=NodeType.END),
        }
        self.edges: list[Edge] = []

        # Create a linear chain of nodes
        previous = "start"
        for i in range(node_count):
            node_name = f"node_{i}"
            self.nodes[node_name] = NodeConfig(
                name=node_name, type=NodeType.FUNCTION, function=f"fn_{i}"
            )
            self.edges.append(Edge(source=previous, target=node_name))
            previous = node_name

        self.edges.append(Edge(source=previous, target="end"))

        # Create function registry with identity functions
        self.fn_registry: dict[str, Any] = {
            f"fn_{i}": lambda x, i=i: x for i in range(node_count)
        }

        self.graph = PureGraph(
            name=f"bench_graph_{node_count}", nodes=self.nodes, edges=self.edges
        )

    def time_topological_order(self, node_count: int) -> None:
        """Benchmark topological ordering."""
        self.graph.topological_order()

    def time_execute(self, node_count: int) -> None:
        """Benchmark graph execution."""
        self.graph.execute(self.fn_registry, initial=0)

    def time_execute_with_result(self, node_count: int) -> None:
        """Benchmark graph execution with result accumulation."""
        result = self.graph.execute(self.fn_registry, initial=42)
        assert result == 42, f"Expected 42, got {result}"