from __future__ import annotations


class BudgetExceededError(RuntimeError):
    pass


def enforce_prompt(estimate: int, plan: dict[str, object]) -> None:
    raw_limit = plan.get("prompt_token_budget")
    if raw_limit is None:
        return
    if not isinstance(raw_limit, (int, str)):
        raise TypeError(f"prompt_token_budget must be int or str, got {type(raw_limit).__name__}")
    limit = int(raw_limit)
    if estimate > limit:
        raise BudgetExceededError(f"prompt estimate {estimate} exceeds plan limit {limit}")


def enforce_requested_completion(requested_tokens: int | None, limit: int) -> None:
    if requested_tokens is None:
        return
    max_tokens = int(requested_tokens)
    if max_tokens > limit:
        raise BudgetExceededError(f"requested completion {max_tokens} exceeds plan limit {limit}")
