from __future__ import annotations

from dataclasses import dataclass, field


class BudgetExceeded(RuntimeError):
    pass


@dataclass
class CostGuard:
    provider_limits: dict[str, float]
    total_limit: float
    pricing_per_million: dict[str, tuple[float, float]]
    provider_spend: dict[str, float] = field(default_factory=dict)
    total_spend: float = 0.0

    def estimate(self, provider: str, input_tokens: int, output_tokens: int) -> float:
        if input_tokens < 0 or output_tokens < 0:
            raise ValueError("Token counts cannot be negative")
        input_price, output_price = self.pricing_per_million[provider]
        return (input_tokens * input_price + output_tokens * output_price) / 1_000_000

    def preflight(self, provider: str, input_tokens: int, output_tokens: int) -> float:
        estimate = self.estimate(provider, input_tokens, output_tokens)
        current = self.provider_spend.get(provider, 0.0)
        limit = self.provider_limits.get(provider, 0.0)
        if limit <= 0 and estimate > 0:
            raise BudgetExceeded(f"{provider} has no positive spend ceiling")
        if current + estimate > limit + 1e-12:
            raise BudgetExceeded(f"{provider} request would exceed its hard ceiling")
        if self.total_spend + estimate > self.total_limit + 1e-12:
            raise BudgetExceeded("Request would exceed the total hard ceiling")
        return estimate

    def record(self, provider: str, input_tokens: int, output_tokens: int) -> float:
        actual = self.estimate(provider, input_tokens, output_tokens)
        self.preflight(provider, input_tokens, output_tokens)
        self.provider_spend[provider] = self.provider_spend.get(provider, 0.0) + actual
        self.total_spend += actual
        return actual

