from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


def _mapping(value: Any) -> dict[str, Any]:
    if value is None:
        return {}
    if isinstance(value, dict):
        return value
    if hasattr(value, "model_dump"):
        return value.model_dump()
    return {
        name: getattr(value, name)
        for name in (
            "prompt_tokens",
            "completion_tokens",
            "total_tokens",
            "input_tokens",
            "output_tokens",
            "cost",
            "prompt_tokens_details",
            "completion_tokens_details",
            "input_tokens_details",
            "output_tokens_details",
        )
        if hasattr(value, name)
    }


def _integer(value: Any) -> int:
    try:
        return int(value or 0)
    except (TypeError, ValueError):
        return 0


def _number(value: Any) -> float | None:
    try:
        return None if value is None else float(value)
    except (TypeError, ValueError):
        return None


@dataclass(slots=True)
class ModelUsage:
    model: str
    calls: int = 0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    reasoning_tokens: int = 0
    cached_tokens: int = 0
    cost: float = 0.0
    unpriced_calls: int = 0
    phases: dict[str, int] = field(default_factory=dict)

    def wire(self) -> dict[str, Any]:
        return {
            "model": self.model,
            "calls": self.calls,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_tokens": self.total_tokens,
            "reasoning_tokens": self.reasoning_tokens,
            "cached_tokens": self.cached_tokens,
            "cost": round(self.cost, 12),
            "unpriced_calls": self.unpriced_calls,
            "phases": dict(sorted(self.phases.items())),
        }


@dataclass(slots=True)
class UsageLedger:
    query_id: str
    by_model: dict[str, ModelUsage] = field(default_factory=dict)

    def record(self, model: str, phase: str, usage: Any) -> None:
        data = _mapping(usage)
        prompt = _integer(data.get("prompt_tokens", data.get("input_tokens")))
        completion = _integer(data.get("completion_tokens", data.get("output_tokens")))
        total = _integer(data.get("total_tokens")) or prompt + completion
        prompt_details = _mapping(
            data.get("prompt_tokens_details", data.get("input_tokens_details"))
        )
        completion_details = _mapping(
            data.get("completion_tokens_details", data.get("output_tokens_details"))
        )
        cached = _integer(prompt_details.get("cached_tokens"))
        reasoning = _integer(completion_details.get("reasoning_tokens"))
        if not reasoning:
            reasoning = _integer(data.get("reasoning_tokens"))
        cost = _number(data.get("cost"))

        item = self.by_model.setdefault(model, ModelUsage(model))
        item.calls += 1
        item.prompt_tokens += prompt
        item.completion_tokens += completion
        item.total_tokens += total
        item.cached_tokens += cached
        item.reasoning_tokens += reasoning
        item.phases[phase] = item.phases.get(phase, 0) + 1
        if cost is None:
            item.unpriced_calls += 1
        else:
            item.cost += cost

    def wire(self) -> dict[str, Any]:
        models = [item.wire() for item in sorted(self.by_model.values(), key=lambda x: x.model)]
        return {
            "query_id": self.query_id,
            "models": models,
            "calls": sum(item.calls for item in self.by_model.values()),
            "prompt_tokens": sum(item.prompt_tokens for item in self.by_model.values()),
            "completion_tokens": sum(item.completion_tokens for item in self.by_model.values()),
            "total_tokens": sum(item.total_tokens for item in self.by_model.values()),
            "reasoning_tokens": sum(item.reasoning_tokens for item in self.by_model.values()),
            "cached_tokens": sum(item.cached_tokens for item in self.by_model.values()),
            "cost": round(sum(item.cost for item in self.by_model.values()), 12),
            "unpriced_calls": sum(item.unpriced_calls for item in self.by_model.values()),
        }