import json
import os
from typing import Optional


class PricingDB:
    """Manages tool and LLM pricing data."""

    def __init__(self, overrides: Optional[dict] = None):
        self._llm_prices: dict = {}
        self._tool_prices: dict = {}
        self._load_defaults()
        if overrides:
            self._apply_overrides(overrides)

    def _load_defaults(self) -> None:
        """Load default pricing from bundled JSON file."""
        pricing_path = os.path.join(
            os.path.dirname(__file__),
            "default_pricing.json",
        )
        try:
            with open(pricing_path) as f:
                data = json.load(f)
                self._llm_prices = data.get("llm", {})
                self._tool_prices = data.get("tools", {})
        except FileNotFoundError:
            pass

    def _apply_overrides(self, overrides: dict) -> None:
        """Apply user-provided pricing overrides."""
        for key, value in overrides.items():
            if isinstance(value, dict) and "input" in value:
                self._llm_prices[key] = value
            elif isinstance(value, (int, float)):
                self._tool_prices[key] = {"cost_per_call": value}
            elif isinstance(value, dict):
                self._tool_prices[key] = value

    def llm_cost(
        self, model: str, tokens_in: int, tokens_out: int
    ) -> float:
        """Calculate cost for an LLM call."""
        # Try exact match first, then prefix match
        pricing = self._llm_prices.get(model)
        if not pricing:
            for key in self._llm_prices:
                if model.startswith(key) or key.startswith(model):
                    pricing = self._llm_prices[key]
                    break

        if not pricing:
            return 0.0

        input_cost = (tokens_in / 1_000_000) * pricing.get("input", 0)
        output_cost = (tokens_out / 1_000_000) * pricing.get("output", 0)
        return round(input_cost + output_cost, 8)

    def tool_cost(self, tool_name: str) -> float:
        """Get cost for a tool call."""
        pricing = self._tool_prices.get(tool_name, {})
        return pricing.get("cost_per_call", 0.0)

    def set_tool_price(self, tool_name: str, cost_per_call: float) -> None:
        """Set or override a tool's price."""
        self._tool_prices[tool_name] = {"cost_per_call": cost_per_call}

    def set_llm_price(
        self, model: str, input_per_mtok: float, output_per_mtok: float
    ) -> None:
        """Set or override an LLM's price per million tokens."""
        self._llm_prices[model] = {
            "input": input_per_mtok,
            "output": output_per_mtok,
        }
