"""Single source of truth for model pricing (USD per 1M tokens).

Prices are loaded from a bundled ``default_pricing.toml`` and optionally
merged with a user-supplied override file.
"""

from __future__ import annotations

import sys
from pathlib import Path

if sys.version_info >= (3, 11):
    import tomllib
else:
    try:
        import tomllib
    except ModuleNotFoundError:  # Python 3.10
        import tomli as tomllib  # type: ignore[no-redef]

_DEFAULT_TOML = Path(__file__).resolve().parent / "default_pricing.toml"

# Module-level cache: list[(pattern, input_cost, output_cost)]
_pricing_cache: list[tuple[str, float, float]] | None = None
_user_pricing_path: str | None = None


def set_user_pricing(path: str | None) -> None:
    """Set (or clear) the path to a user override TOML and invalidate the cache."""
    global _user_pricing_path, _pricing_cache
    _user_pricing_path = path
    _pricing_cache = None


def _load_toml(path: Path) -> dict[str, dict[str, float]]:
    """Load a pricing TOML and return the ``[models]`` table."""
    with open(path, "rb") as f:
        data = tomllib.load(f)
    return data.get("models", {})


def _load_pricing() -> list[tuple[str, float, float]]:
    """Load default pricing, merge user overrides, return ordered list."""
    global _pricing_cache
    if _pricing_cache is not None:
        return _pricing_cache

    models = _load_toml(_DEFAULT_TOML)

    # User override: explicit path > auto-detect cwd file
    override_path = _user_pricing_path
    if override_path is None:
        cwd_candidate = Path.cwd() / "minitrail_pricing.toml"
        if cwd_candidate.is_file():
            override_path = str(cwd_candidate)

    if override_path:
        user_models = _load_toml(Path(override_path))
        models.update(user_models)  # user values win

    _pricing_cache = [
        (name, entry["input"], entry["output"])
        for name, entry in models.items()
    ]
    return _pricing_cache


def cost_per_1m(model: str) -> tuple[float, float] | None:
    """Return (input, output) cost per 1M tokens, or None if model unknown."""
    low = model.lower()
    for pat, pi, po in _load_pricing():
        if pat in low:
            return (pi, po)
    return None


def compute_cost(model: str, inp: int, out: int) -> float:
    """Compute cost in USD for given token counts."""
    pricing = cost_per_1m(model)
    if pricing is None:
        return 0.0
    pi, po = pricing
    return (inp / 1_000_000) * pi + (out / 1_000_000) * po
