from __future__ import annotations

from dataclasses import dataclass
from typing import Optional
import re

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


@dataclass(frozen=True)
class ModelKey:
    provider: str
    model: str

_DATE_SUFFIX = re.compile(r"-\d{4}-\d{2}-\d{2}$")

def canonicalize(provider: str, model: str) -> ModelKey:
    p = provider.strip().lower()

    m = model.strip()
    # ensure "openai/..." prefix is present if that's your convention
    if "/" not in m:
        m = f"{p}/{m}"

    # strip dated suffix: openai/gpt-4.1-2025-04-14 -> openai/gpt-4.1
    m = _DATE_SUFFIX.sub("", m)

    return ModelKey(provider=p, model=m)


class PricingRegistry:
    def __init__(self):
        self._prices: dict[ModelKey, ModelPrice] = {}

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

    def get_price(self, provider: str, model: str) -> ModelPrice:
        key = canonicalize(provider, model)
        try:
            return self._prices[key]
        except KeyError as e:
            raise KeyError(f"No pricing configured for {key.provider}/{key.model}") from e