"""Cost estimation over model rates and token counts."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from llm_connect.fx import FxRate, resolve_fx_rate
from llm_connect.rates import ModelRateRegistry


@dataclass(frozen=True)
class CostEstimate:
    """Cost estimate split by prompt and completion token spend.

    USD fields come from the rate table. EUR fields are derived via :class:`FxRate`
    when conversion is available; missing FX yields ``cost_eur=None`` with an
    explicit ``fx_source`` rather than silently treating spend as zero.
    """

    cost_usd: float | None
    cost_source: str
    prompt_cost_usd: float | None = None
    completion_cost_usd: float | None = None
    cost_eur: float | None = None
    prompt_cost_eur: float | None = None
    completion_cost_eur: float | None = None
    fx_source: str | None = None


def estimate_cost(
    model_id: str,
    prompt_tokens: int,
    completion_tokens: int = 0,
    *,
    registry: ModelRateRegistry | None = None,
    fx: FxRate | float | None = None,
    apply_fx: bool = True,
) -> CostEstimate:
    """Estimate USD (and optionally EUR) cost for token counts using *registry*.

    Unknown models return ``CostEstimate(None, "unknown")`` so callers can
    record uncertainty explicitly instead of treating missing prices as zero.

    When *apply_fx* is true (default), EUR fields are filled using *fx* or the
    resolved default FX rate. Pass ``apply_fx=False`` to skip conversion.
    """
    prompt_count = _non_negative_int("prompt_tokens", prompt_tokens)
    completion_count = _non_negative_int("completion_tokens", completion_tokens)
    rates = registry or ModelRateRegistry.default()
    rate = rates.get(model_id)
    if rate is None:
        return CostEstimate(cost_usd=None, cost_source="unknown")

    prompt_cost = (prompt_count / 1000.0) * rate.prompt_per_1k
    completion_cost = (completion_count / 1000.0) * rate.completion_per_1k
    cost_usd = prompt_cost + completion_cost

    cost_eur = None
    prompt_cost_eur = None
    completion_cost_eur = None
    fx_source: str | None = None
    if apply_fx:
        resolved = resolve_fx_rate(fx)
        if resolved is None:
            fx_source = "unknown"
        else:
            fx_source = resolved.source
            cost_eur = resolved.usd_to_eur(cost_usd)
            prompt_cost_eur = resolved.usd_to_eur(prompt_cost)
            completion_cost_eur = resolved.usd_to_eur(completion_cost)

    return CostEstimate(
        cost_usd=cost_usd,
        cost_source=f"rate_table:{rate.model_id}",
        prompt_cost_usd=prompt_cost,
        completion_cost_usd=completion_cost,
        cost_eur=cost_eur,
        prompt_cost_eur=prompt_cost_eur,
        completion_cost_eur=completion_cost_eur,
        fx_source=fx_source,
    )


@dataclass(frozen=True)
class CostModel:
    """Small wrapper for callers that prefer an object over a free function."""

    registry: ModelRateRegistry | None = None
    fx: FxRate | float | None = None
    apply_fx: bool = True

    def estimate_cost(
        self,
        model_id: str,
        prompt_tokens: int,
        completion_tokens: int = 0,
    ) -> CostEstimate:
        """Estimate cost using this model's registry and FX settings."""
        return estimate_cost(
            model_id,
            prompt_tokens,
            completion_tokens,
            registry=self.registry,
            fx=self.fx,
            apply_fx=self.apply_fx,
        )


def _non_negative_int(name: str, value: Any) -> int:
    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        raise ValueError(f"{name} must be a non-negative integer")
    return value