"""Benchmark runner with InstrumentedProvider for capturing metrics."""

from __future__ import annotations

import logging
import time
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Any

from sage.models import CompletionResult, Message, StreamChunk, ToolSchema, Usage
from sage.providers.litellm_provider import LiteLLMProvider

from sage_evaluator.models import ModelRunResult, ToolCallTrace

if TYPE_CHECKING:
    from sage.main_config import MainConfig

logger = logging.getLogger(__name__)


class InstrumentedProvider(LiteLLMProvider):
    """LiteLLMProvider wrapper that accumulates token counts and tool call traces.

    Injected into Agent(provider=...) to capture per-call metrics
    without modifying apollo-agent-v2.
    """

    def __init__(self, model: str, **kwargs: Any) -> None:
        super().__init__(model, **kwargs)
        self.total_prompt_tokens: int = 0
        self.total_completion_tokens: int = 0
        self.total_tokens: int = 0
        self.call_count: int = 0
        self.tool_call_traces: list[ToolCallTrace] = []
        self._start_time: float = 0.0
        self._total_latency_ms: float = 0.0

    def reset(self) -> None:
        """Reset all accumulated metrics."""
        self.total_prompt_tokens = 0
        self.total_completion_tokens = 0
        self.total_tokens = 0
        self.call_count = 0
        self.tool_call_traces = []
        self._start_time = 0.0
        self._total_latency_ms = 0.0

    def start_timer(self) -> None:
        """Start the overall run timer."""
        self._start_time = time.perf_counter()

    @property
    def elapsed_ms(self) -> float:
        """Total elapsed time since start_timer() in milliseconds."""
        if self._start_time == 0.0:
            return 0.0
        return (time.perf_counter() - self._start_time) * 1000

    async def complete(
        self,
        messages: list[Message],
        tools: list[ToolSchema] | None = None,
        **kwargs: Any,
    ) -> CompletionResult:
        """Completion with metrics capture."""
        call_start = time.perf_counter()
        result = await super().complete(messages, tools, **kwargs)
        call_duration_ms = (time.perf_counter() - call_start) * 1000

        self._accumulate_usage(result.usage)
        self.call_count += 1
        self._total_latency_ms += call_duration_ms

        # Trace tool calls if present
        if result.message.tool_calls:
            for tc in result.message.tool_calls:
                self.tool_call_traces.append(
                    ToolCallTrace(
                        name=tc.name,
                        arguments=tc.arguments,
                        duration_ms=call_duration_ms,
                    )
                )

        return result

    async def stream(
        self,
        messages: list[Message],
        tools: list[ToolSchema] | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[StreamChunk]:
        """Streaming completion with metrics capture."""
        call_start = time.perf_counter()
        last_usage: Usage | None = None

        async for chunk in super().stream(messages, tools, **kwargs):
            if chunk.usage is not None:
                last_usage = chunk.usage
            if chunk.tool_calls:
                for tc in chunk.tool_calls:
                    self.tool_call_traces.append(
                        ToolCallTrace(
                            name=tc.name,
                            arguments=tc.arguments,
                        )
                    )
            yield chunk

        call_duration_ms = (time.perf_counter() - call_start) * 1000
        self._total_latency_ms += call_duration_ms
        self.call_count += 1

        if last_usage:
            self._accumulate_usage(last_usage)

    def _accumulate_usage(self, usage: Usage) -> None:
        """Add usage stats to running totals."""
        self.total_prompt_tokens += usage.prompt_tokens
        self.total_completion_tokens += usage.completion_tokens
        self.total_tokens += usage.total_tokens

    def to_run_result(self, output: str, error: str | None = None) -> ModelRunResult:
        """Convert accumulated metrics into a ModelRunResult."""
        return ModelRunResult(
            model=self.model,
            output=output,
            prompt_tokens=self.total_prompt_tokens,
            completion_tokens=self.total_completion_tokens,
            total_tokens=self.total_tokens,
            latency_ms=self.elapsed_ms,
            tool_calls=list(self.tool_call_traces),
            error=error,
            success=error is None,
        )


class BenchmarkRunner:
    """Runs a single agent execution against a specific model and captures metrics."""

    async def run(
        self,
        config_path: str,
        model: str,
        intent: str,
        model_params: dict[str, Any] | None = None,
        central: MainConfig | None = None,
    ) -> ModelRunResult:
        """Run an agent with an instrumented provider and return metrics.

        Args:
            config_path: Path to agent config (AGENTS.md or directory).
            model: LiteLLM model identifier to use.
            intent: User input to send to the agent.
            model_params: Optional model parameters (temperature, etc.).
            central: Optional main config for applying config.toml defaults.

        Returns:
            ModelRunResult with token counts, latency, and tool call traces.
        """
        from sage.agent import Agent
        from sage.config import load_config

        logger.info("Loading agent config from %s", config_path)
        config = load_config(config_path, central=central)
        provider = InstrumentedProvider(model, **(model_params or {}))

        # Build agent from config but override the model and provider
        agent = Agent(
            name=config.name or "benchmark-agent",
            model=model,
            description=config.description or "",
            tools=[t for t in (config.tools or [])],
            provider=provider,
            max_turns=config.max_turns,
            body=config._body or "",
            model_params=config.model_params.to_kwargs() if config.model_params else None,
        )

        provider.start_timer()
        logger.info("Executing agent with model %s", model)

        try:
            output = await agent.run(intent)
            logger.info(
                "Run complete for model %s: %d tokens, %.0fms latency",
                model,
                provider.total_tokens,
                provider.elapsed_ms,
            )
            return provider.to_run_result(output=output)
        except Exception as e:
            logger.error("Benchmark run failed for model %s: %s", model, e)
            return provider.to_run_result(output="", error=str(e))
