from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ember.models.pricing import Pricing as RegistryPricingImpl


@dataclass(frozen=True)
class PricingDimensions:
    """Inputs required to price a request or chunk."""

    model: str
    prompt_tokens: int
    completion_tokens: int
    price_book_version: str | None = None
    billing_mode: str = "resell"

    @property
    def total_tokens(self) -> int:
        return max(0, self.prompt_tokens) + max(0, self.completion_tokens)


class Pricing(ABC):
    """Pricing abstraction used across the gateway."""

    @abstractmethod
    def calculate(self, dims: PricingDimensions) -> float:
        """Return the USD cost for the provided token usage."""


class NaivePricing(Pricing):
    """Fallback pricing using a flat $0.01 / 1K token heuristic."""

    def calculate(self, dims: PricingDimensions) -> float:
        if dims.total_tokens == 0:
            return 0.0
        return round((dims.total_tokens / 1000.0) * 0.01, 6)


class RegistryPricing(Pricing):
    """Wrapper around the shared pricing registry with lazy import."""

    _pricing: RegistryPricingImpl | None

    def __init__(self) -> None:
        self._pricing = None

    def _ensure_impl(self) -> RegistryPricingImpl:
        if self._pricing is None:
            from ember.models.pricing import Pricing as RegistryPricingImpl

            self._pricing = RegistryPricingImpl()
        return self._pricing

    def calculate(self, dims: PricingDimensions) -> float:
        if dims.billing_mode.lower() == "byok":
            return 0.0

        impl = self._ensure_impl()
        cost = impl.calculate_cost(
            dims.model,
            input_tokens=max(0, dims.prompt_tokens),
            output_tokens=max(0, dims.completion_tokens),
        )
        return round(float(cost), 6)


class CompositePricing(Pricing):
    """Primary pricing strategy with registry first, naive fallback."""

    def __init__(self, primary: Pricing | None = None, fallback: Pricing | None = None) -> None:
        self._primary = primary or RegistryPricing()
        self._fallback = fallback or NaivePricing()

    def calculate(self, dims: PricingDimensions) -> float:
        cost = self._primary.calculate(dims)
        if cost == 0.0 and dims.total_tokens > 0 and dims.billing_mode.lower() != "byok":
            return self._fallback.calculate(dims)
        return cost


def default_pricing() -> Pricing:
    """Return the default pricing strategy for the gateway."""

    return CompositePricing()


__all__ = [
    "PricingDimensions",
    "Pricing",
    "NaivePricing",
    "RegistryPricing",
    "CompositePricing",
    "default_pricing",
]
