"""Pricing sources and cost calculation for nanomon."""

from __future__ import annotations

import json
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Callable

from nanomon.errors import PricingError
from nanomon.types import ModelPrice


class PricingSource(ABC):
    """Abstract base class for pricing sources."""

    @abstractmethod
    def get_prices(self, model: str) -> ModelPrice | None:
        """Get pricing for a model.

        Args:
            model: The model name to look up

        Returns:
            ModelPrice if found, None otherwise
        """
        ...


class DictPricing(PricingSource):
    """Pricing source backed by an in-memory dictionary.

    Expected format:
        {
            "model-name": {"input_per_1m": 1.0, "output_per_1m": 2.0},
            ...
        }
    """

    def __init__(self, prices: dict[str, dict[str, float]]) -> None:
        """Initialize with a pricing dictionary.

        Args:
            prices: Dictionary mapping model names to pricing info
        """
        self._prices = prices

    def get_prices(self, model: str) -> ModelPrice | None:
        """Get pricing for a model."""
        if model not in self._prices:
            return None

        data = self._prices[model]
        return ModelPrice(
            input_per_1m=data.get("input_per_1m", 0.0),
            output_per_1m=data.get("output_per_1m", 0.0),
        )


class JsonPricing(PricingSource):
    """Pricing source backed by a JSON file.

    Expected JSON format:
        {
            "model-name": {"input_per_1m": 1.0, "output_per_1m": 2.0},
            ...
        }
    """

    def __init__(self, path: str | Path) -> None:
        """Initialize with a JSON file path.

        Args:
            path: Path to the JSON pricing file

        Raises:
            PricingError: If the file cannot be read or parsed
        """
        self._path = Path(path)
        self._prices: dict[str, dict[str, float]] = {}
        self._load()

    def _load(self) -> None:
        """Load pricing data from the JSON file."""
        try:
            with self._path.open("r", encoding="utf-8") as f:
                self._prices = json.load(f)
        except FileNotFoundError as e:
            raise PricingError(f"Pricing file not found: {self._path}") from e
        except json.JSONDecodeError as e:
            raise PricingError(f"Invalid JSON in pricing file: {self._path}") from e

    def get_prices(self, model: str) -> ModelPrice | None:
        """Get pricing for a model."""
        if model not in self._prices:
            return None

        data = self._prices[model]
        return ModelPrice(
            input_per_1m=data.get("input_per_1m", 0.0),
            output_per_1m=data.get("output_per_1m", 0.0),
        )


class CallablePricing(PricingSource):
    """Pricing source backed by a callable resolver.

    The callable should accept a model name and return a ModelPrice or None.
    """

    def __init__(self, resolver: Callable[[str], ModelPrice | None]) -> None:
        """Initialize with a callable resolver.

        Args:
            resolver: Function that takes model name and returns ModelPrice or None
        """
        self._resolver = resolver

    def get_prices(self, model: str) -> ModelPrice | None:
        """Get pricing for a model."""
        return self._resolver(model)


def calculate_cost(
    input_tokens: int,
    output_tokens: int,
    price: ModelPrice | None,
) -> float:
    """Calculate the cost for a call given token counts and pricing.

    Args:
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        price: Pricing information (or None if unknown)

    Returns:
        Cost in USD, 0.0 if pricing is unknown
    """
    if price is None:
        return 0.0

    input_cost = (input_tokens / 1_000_000) * price.input_per_1m
    output_cost = (output_tokens / 1_000_000) * price.output_per_1m
    return input_cost + output_cost
