from __future__ import annotations

from decimal import Decimal
from typing import Any, Mapping


def _to_non_negative_int(value: Any, *, field: str) -> int:
    if isinstance(value, bool):
        raise ValueError(f"{field} must be an integer")
    if isinstance(value, int):
        converted = value
    elif isinstance(value, str) and value.isdigit():
        converted = int(value)
    else:
        raise ValueError(f"{field} must be an integer")
    if converted < 0:
        raise ValueError(f"{field} cannot be negative")
    return converted


def _to_non_negative_decimal(value: Any, *, field: str) -> Decimal:
    converted = Decimal(str(value))
    if converted < 0:
        raise ValueError(f"{field} cannot be negative")
    return converted


class CostTracker:
    def __init__(self, *, currency: str = "USD") -> None:
        self._input_tokens = 0
        self._output_tokens = 0
        self._total_tokens = 0
        self._cost = Decimal("0")
        self._currency = currency

    def add_usage(self, usage: Mapping[str, Any]) -> None:
        input_tokens = usage.get("input_tokens", usage.get("prompt_tokens", 0))
        output_tokens = usage.get("output_tokens", usage.get("completion_tokens", 0))
        total_tokens = usage.get("total_tokens")
        cost = usage.get("cost", "0")
        currency = usage.get("currency")

        normalized_input = _to_non_negative_int(input_tokens, field="input_tokens")
        normalized_output = _to_non_negative_int(output_tokens, field="output_tokens")
        normalized_total = (
            _to_non_negative_int(total_tokens, field="total_tokens")
            if total_tokens is not None
            else normalized_input + normalized_output
        )
        normalized_cost = _to_non_negative_decimal(cost, field="cost")

        self._input_tokens += normalized_input
        self._output_tokens += normalized_output
        self._total_tokens += normalized_total
        self._cost += normalized_cost

        if currency is not None:
            normalized_currency = str(currency)
            if normalized_currency != self._currency:
                raise ValueError("currency mismatch")

    def snapshot(self) -> dict[str, Any]:
        return {
            "input_tokens": self._input_tokens,
            "output_tokens": self._output_tokens,
            "total_tokens": self._total_tokens,
            "cost": self._cost,
            "currency": self._currency,
        }
