from __future__ import annotations

from dataclasses import dataclass
from typing import Optional


@dataclass(frozen=True)
class ModelPrice:
    input_per_million: float
    output_per_million: float


class PricingRegistry:
    """
    Simple in-memory pricing table.
    Users can override by calling `set_price`.
    """

    def __init__(self) -> None:
        self._prices: dict[tuple[str, str], ModelPrice] = {}

    def set_price(self, provider: str, model: str, *, input_per_million: float, output_per_million: float) -> None:
        self._prices[(provider.lower(), model)] = ModelPrice(
            input_per_million=float(input_per_million),
            output_per_million=float(output_per_million),
        )

    def get_price(self, provider: str, model: str) -> Optional[ModelPrice]:
        return self._prices.get((provider.lower(), model))
