"""Cost tracking for Logos ADK agents."""
from __future__ import annotations

from dataclasses import dataclass
from fnmatch import fnmatch

from logos_adk.observe import Exporter, Metric, Span


# Prices per 1M tokens: (input_usd, output_usd)
DEFAULT_RATES: dict[str, tuple[float, float]] = {
    "claude-opus-4*": (15.0, 75.0),
    "claude-sonnet-4*": (3.0, 15.0),
    "claude-haiku-4*": (0.80, 4.0),
    "claude-3-5-sonnet*": (3.0, 15.0),
    "claude-3-5-haiku*": (0.80, 4.0),
    "claude-3-opus*": (15.0, 75.0),
    "claude-3-sonnet*": (3.0, 15.0),
    "claude-3-haiku*": (0.25, 1.25),
    "gpt-4o*": (2.50, 10.0),
    "gpt-4o-mini*": (0.15, 0.60),
    "gpt-4-turbo*": (10.0, 30.0),
    "gpt-4": (30.0, 60.0),
    "gpt-3.5-turbo*": (0.50, 1.50),
    "o1*": (15.0, 60.0),
    "o3*": (10.0, 40.0),
}


@dataclass
class CostReport:
    """Aggregated cost report."""

    total_usd: float = 0.0
    input_usd: float = 0.0
    output_usd: float = 0.0
    input_tokens: int = 0
    output_tokens: int = 0
    calls: int = 0


@dataclass
class _CostEntry:
    agent_name: str | None
    session_id: str | None
    model: str
    input_tokens: int
    output_tokens: int
    input_usd: float
    output_usd: float


class CostTracker(Exporter):
    """Tracks token costs by intercepting LLM spans."""

    def __init__(self, rates: dict[str, tuple[float, float]] | None = None) -> None:
        self._rates: dict[str, tuple[float, float]] = dict(DEFAULT_RATES)
        if rates:
            self._rates.update(rates)
        self._entries: list[_CostEntry] = []

    def _resolve_rate(self, model: str) -> tuple[float, float]:
        """Find rate for model via exact or glob match."""
        if model in self._rates:
            return self._rates[model]
        for pattern, rate in self._rates.items():
            if "*" in pattern and fnmatch(model, pattern):
                return rate
        return (0.0, 0.0)

    async def on_span(self, span: Span) -> None:
        if span.name not in ("LLM.generate", "LLM.stream"):
            return

        input_tokens = span.attributes.get("input_tokens", 0)
        output_tokens = span.attributes.get("output_tokens", 0)
        model = span.attributes.get("model", "unknown")

        input_rate, output_rate = self._resolve_rate(model)
        input_usd = input_tokens * input_rate / 1_000_000
        output_usd = output_tokens * output_rate / 1_000_000

        self._entries.append(_CostEntry(
            agent_name=span.agent_name,
            session_id=span.session_id,
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            input_usd=input_usd,
            output_usd=output_usd,
        ))

    async def on_metric(self, metric: Metric) -> None:
        pass

    async def flush(self) -> None:
        pass

    @property
    def total(self) -> CostReport:
        """Get total cost report across all agents and sessions."""
        return self.report()

    def report(
        self,
        *,
        agent_name: str | None = None,
        session_id: str | None = None,
    ) -> CostReport:
        """Get filtered cost report."""
        entries = self._entries
        if agent_name is not None:
            entries = [e for e in entries if e.agent_name == agent_name]
        if session_id is not None:
            entries = [e for e in entries if e.session_id == session_id]

        return CostReport(
            total_usd=sum(e.input_usd + e.output_usd for e in entries),
            input_usd=sum(e.input_usd for e in entries),
            output_usd=sum(e.output_usd for e in entries),
            input_tokens=sum(e.input_tokens for e in entries),
            output_tokens=sum(e.output_tokens for e in entries),
            calls=len(entries),
        )

    def reset(self) -> None:
        """Clear all tracked cost data."""
        self._entries.clear()
