"""Airspeed Velocity benchmarks for plan generation workflows.

Measures invoke and streaming performance of PlanGenerationGraph using
built-in FakeListLLM responses and minimal domain models.
"""

from __future__ import annotations

from pathlib import Path

import tiktoken
from langchain_community.llms import FakeListLLM

from cleveragents.agents.graphs.plan_generation import PlanGenerationGraph
from cleveragents.domain.models.core import (
    Context,
    ContextType,
    Plan,
    PlanStatus,
    Project,
)
from cleveragents.domain.models.core.project_legacy import ProjectSettings


def _sample_project() -> Project:
    return Project(
        id=1,
        name="bench-project",
        path=Path(".").resolve(),
        settings=ProjectSettings(
            auto_build=False,
            auto_apply=False,
            confirm_apply=True,
            max_context_size=52_428_800,
            default_model="mock-gpt",
        ),
        current_plan_id=1,
    )


def _sample_plan() -> Plan:
    return Plan(
        id=1,
        project_id=1,
        name="bench-plan",
        prompt="Add error handling",
        status=PlanStatus.PENDING,
        current=True,
        build=None,
        build_started_at=None,
        build_completed_at=None,
        model_used=None,
        token_count=None,
        result=None,
        applied_at=None,
        files_created=None,
        files_modified=None,
        files_deleted=None,
    )


def _sample_contexts() -> list[Context]:
    return [
        Context(
            id=1,
            plan_id=1,
            type=ContextType.FILE,
            path="src/example.py",
            content="print('hello')",
            file_hash=None,
            size=18,
        )
    ]


class PlanGenerationSuite:
    """Benchmark PlanGenerationGraph invoke and stream paths."""

    def setup(self) -> None:
        self.graph = PlanGenerationGraph(
            llm=FakeListLLM(responses=["bench response"] * 3)
        )
        self.project = _sample_project()
        self.plan = _sample_plan()
        self.contexts = _sample_contexts()

        enc = tiktoken.get_encoding("cl100k_base")
        self.prompt_tokens = {
            "analyze": len(enc.encode(self.graph.analyze_prompt.template)),
            "generate": len(enc.encode(self.graph.generate_prompt.template)),
            "validate": len(enc.encode(self.graph.validate_prompt.template)),
        }

    def time_invoke(self) -> None:
        self.graph.invoke(self.project, self.plan, self.contexts, thread_id="bench")

    def time_stream(self) -> None:
        for _ in self.graph.stream(self.project, self.plan, self.contexts):
            pass

    def track_analyze_prompt_tokens(self) -> int:
        return self.prompt_tokens["analyze"]

    def track_generate_prompt_tokens(self) -> int:
        return self.prompt_tokens["generate"]

    def track_validate_prompt_tokens(self) -> int:
        return self.prompt_tokens["validate"]