"""Cost calculator for routing decisions.

Computes per-request costs using pricing data from models_db.json.
Stateless design — all state lives in the models database.

Follows the immutable cost tracking pattern from cost-aware-llm-pipeline skill:
results are frozen Pydantic models, pricing lookups are pure functions.
"""

from __future__ import annotations

import structlog
from pydantic import BaseModel

from robot_resources.routing.selector import MODELS_DB

logger = structlog.get_logger()


class CostResult(BaseModel):
    """Immutable result of a cost calculation for a single request.

    Contains actual cost, baseline cost, and the difference (savings).
    """

    cost_actual: float
    cost_baseline: float
    cost_saved: float
    model_used: str
    model_baseline: str

    model_config = {"frozen": True}

    @property
    def savings_percent(self) -> float:
        """Percentage saved vs baseline. 0.0 if baseline is zero."""
        if self.cost_baseline == 0.0:
            return 0.0
        return ((self.cost_baseline - self.cost_actual) / self.cost_baseline) * 100.0


class CostCalculator:
    """Computes per-request costs from models_db pricing.

    Pricing is read from models_db.json ``cost_per_1k_input`` and
    ``cost_per_1k_output`` fields. An optional custom ``models_db``
    list can be passed for testing.

    Usage::

        calc = CostCalculator()
        cost = calc.calculate_cost("gpt-5.4", input_tokens=1000, output_tokens=500)
        result = calc.calculate_request_costs("gpt-5.4-mini", "gpt-5.4", 1000, 500)
    """

    def __init__(self, models_db: list[dict] | None = None) -> None:
        db = models_db if models_db is not None else MODELS_DB
        # Build name → (cost_per_1k_input, cost_per_1k_output) index
        self._pricing: dict[str, tuple[float, float]] = {}
        for model in db:
            name = model["name"]
            cost_in = model.get("cost_per_1k_input", 0.0)
            cost_out = model.get("cost_per_1k_output", 0.0)
            self._pricing[name] = (cost_in, cost_out)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def get_pricing(self, model_name: str) -> tuple[float, float] | None:
        """Look up pricing for a model.

        Returns ``(cost_per_1k_input, cost_per_1k_output)`` or ``None``
        if the model is not in the database.
        """
        return self._pricing.get(model_name)

    def calculate_cost(
        self,
        model_name: str,
        *,
        input_tokens: int,
        output_tokens: int,
    ) -> float:
        """Calculate USD cost for a single model invocation.

        Returns 0.0 and logs a warning if the model has no pricing data.
        Negative token counts are clamped to zero.
        """
        pricing = self._pricing.get(model_name)
        if pricing is None:
            logger.warning(
                "cost_calculator_unknown_model",
                model=model_name,
                note="Returning 0.0 — model not in pricing database",
            )
            return 0.0

        cost_per_1k_in, cost_per_1k_out = pricing
        safe_in = max(0, input_tokens)
        safe_out = max(0, output_tokens)
        return (safe_in / 1000.0) * cost_per_1k_in + (safe_out / 1000.0) * cost_per_1k_out

    def calculate_request_costs(
        self,
        model_used: str,
        model_baseline: str,
        input_tokens: int,
        output_tokens: int,
    ) -> CostResult:
        """Calculate actual, baseline, and saved costs for a routed request.

        This is the main entry point used by OutcomeRecorder.
        """
        cost_actual = self.calculate_cost(
            model_used, input_tokens=input_tokens, output_tokens=output_tokens
        )
        cost_baseline = self.calculate_cost(
            model_baseline, input_tokens=input_tokens, output_tokens=output_tokens
        )
        return CostResult(
            cost_actual=cost_actual,
            cost_baseline=cost_baseline,
            cost_saved=cost_baseline - cost_actual,
            model_used=model_used,
            model_baseline=model_baseline,
        )
