"""OpenAI token usage and cost metering."""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from threading import Lock
from typing import Any

from app.config import get_settings


@dataclass
class UsageEvent:
    """One LLM API call."""

    kind: str
    model: str
    input_tokens: int
    output_tokens: int
    elapsed_ms: float
    is_vision: bool = False

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens


@dataclass
class UsageSummary:
    """Aggregated usage + USD cost."""

    calls: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    cost_input_usd: float = 0.0
    cost_output_usd: float = 0.0
    cost_total_usd: float = 0.0
    by_model: dict[str, dict[str, Any]] = field(default_factory=dict)
    events: list[dict[str, Any]] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Serialize for JSON reports."""
        return {
            "calls": self.calls,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "total_tokens": self.total_tokens,
            "cost_input_usd": round(self.cost_input_usd, 6),
            "cost_output_usd": round(self.cost_output_usd, 6),
            "cost_total_usd": round(self.cost_total_usd, 6),
            "by_model": self.by_model,
            "events": self.events,
        }


class UsageMeter:
    """Thread-safe accumulator for LLM usage."""

    def __init__(self) -> None:
        self._lock = Lock()
        self._events: list[UsageEvent] = []

    def reset(self) -> None:
        """Clear recorded events."""
        with self._lock:
            self._events.clear()

    def record(self, event: UsageEvent) -> None:
        """Append one usage event."""
        with self._lock:
            self._events.append(event)

    def summary(self) -> UsageSummary:
        """Compute totals and USD costs."""
        settings = get_settings()
        with self._lock:
            events = list(self._events)

        out = UsageSummary()
        for event in events:
            out.calls += 1
            out.input_tokens += event.input_tokens
            out.output_tokens += event.output_tokens
            out.total_tokens += event.total_tokens
            out.events.append(asdict(event))

            if event.is_vision:
                in_rate = settings.openai_vision_input_per_1m
                out_rate = settings.openai_vision_output_per_1m
            else:
                in_rate = settings.openai_text_input_per_1m
                out_rate = settings.openai_text_output_per_1m

            cost_in = event.input_tokens * in_rate / 1_000_000
            cost_out = event.output_tokens * out_rate / 1_000_000
            out.cost_input_usd += cost_in
            out.cost_output_usd += cost_out

            bucket = out.by_model.setdefault(
                event.model,
                {
                    "calls": 0,
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "cost_usd": 0.0,
                    "is_vision": event.is_vision,
                },
            )
            bucket["calls"] += 1
            bucket["input_tokens"] += event.input_tokens
            bucket["output_tokens"] += event.output_tokens
            bucket["cost_usd"] = round(
                bucket["cost_usd"] + cost_in + cost_out, 6
            )

        out.cost_total_usd = out.cost_input_usd + out.cost_output_usd
        for model, bucket in out.by_model.items():
            bucket["cost_usd"] = round(bucket["cost_usd"], 6)
        return out


_METER = UsageMeter()


def get_usage_meter() -> UsageMeter:
    """Return process-wide usage meter."""
    return _METER