"""Budget limit helpers shared across gateway request pipelines."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Mapping

from .errors import GatewayError


@dataclass(frozen=True)
class BudgetLimits:
    prompt_tokens: int
    completion_tokens: int


_DEFAULT_LIMITS = BudgetLimits(prompt_tokens=2_000_000, completion_tokens=2_000_000)


def _coerce_limit(
    plan: Mapping[str, object],
    *,
    field: str,
    provider: str,
) -> int:
    raw_value = plan.get(field)
    if raw_value is None:
        raise GatewayError(
            "plan_missing_budget",
            provider=provider,
            context={"field": field},
        )
    if not isinstance(raw_value, (int, str, float)):
        raise GatewayError(
            "invalid_plan_budget",
            provider=provider,
            context={"field": field, "value": raw_value},
        )
    try:
        value: int = int(raw_value)
    except (TypeError, ValueError) as exc:
        raise GatewayError(
            "invalid_plan_budget",
            provider=provider,
            context={"field": field, "value": raw_value},
        ) from exc
    if value <= 0:
        raise GatewayError(
            "invalid_plan_budget",
            provider=provider,
            context={"field": field, "value": raw_value},
        )
    return value


def resolve_budget_limits(
    plan: Mapping[str, object],
    *,
    provider: str,
) -> BudgetLimits:
    """Return enforced budget limits for the resolved provider."""

    normalized = provider or "unknown"

    prompt_limit = _coerce_limit(plan, field="prompt_token_budget", provider=normalized)
    completion_limit = _coerce_limit(
        plan,
        field="completion_token_budget",
        provider=normalized,
    )

    if normalized.startswith("takumi"):
        return BudgetLimits(prompt_tokens=prompt_limit, completion_tokens=completion_limit)

    return BudgetLimits(
        prompt_tokens=min(prompt_limit, _DEFAULT_LIMITS.prompt_tokens),
        completion_tokens=min(completion_limit, _DEFAULT_LIMITS.completion_tokens),
    )


__all__ = ["BudgetLimits", "resolve_budget_limits"]
