import json
import os
import aiohttp
import re
from typing import Dict
from functools import lru_cache

class Valuator:
    """
    A library to fetch model_prices.json from GitHub and perform regex search on model names,
    returning only input and output cost per token for matched models.
    """

    def __init__(self, model_prices_url: str = "https://raw.githubusercontent.com/AgentOps-AI/tokencost/main/tokencost/model_prices.json", cache_file: str = "model_prices.json"):
        """
        Initialize Valuator with configuration for fetching model prices.

        Args:
            model_prices_url (str): URL to fetch model_prices.json.
            cache_file (str): Path to cache model_prices.json locally.

        Note:
            Call `await initialize()` to complete setup.
        """
        self.model_prices_url = model_prices_url
        self.cache_file = cache_file
        self.model_prices = {}
        self.model_names = set()
        self._session = None

    async def initialize(self):
        """
        Asynchronously initialize by fetching model prices and setting up the session.
        """
        if self._session is None:
            self._session = aiohttp.ClientSession()
            try:
                await self._fetch_model_prices()
            except Exception as e:
                await self._session.close()
                self._session = None
                raise e

    async def _fetch_model_prices(self):
        """
        Fetch model_prices.json from the URL or load from cache if available and valid.
        Only load necessary fields to reduce memory usage.
        """
        if os.path.exists(self.cache_file):
            try:
                with open(self.cache_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    self.model_prices = {
                        key: {
                            "input_cost_per_token": value.get("input_cost_per_token", 0.0),
                            "output_cost_per_token": value.get("output_cost_per_token", 0.0)
                        }
                        for key, value in data.items()
                    }
                    self.model_names = set(self.model_prices.keys())
                return
            except (json.JSONDecodeError, IOError):
                print("Invalid or inaccessible cache file, fetching from URL...")

        async with self._session.get(self.model_prices_url) as response:
            if response.status == 200:
                response_text = await response.text()
                data = json.loads(response_text)
                self.model_prices = {
                    key: {
                        "input_cost_per_token": value.get("input_cost_per_token", 0.0),
                        "output_cost_per_token": value.get("output_cost_per_token", 0.0)
                    }
                    for key, value in data.items()
                }
                self.model_names = set(self.model_prices.keys())
                try:
                    with open(self.cache_file, 'w', encoding='utf-8') as f:
                        json.dump(self.model_prices, f, indent=2)
                except IOError:
                    print("Failed to cache model prices locally.")
            else:
                raise Exception(f"Failed to fetch model prices: {response.status}")

    async def close(self):
        """
        Close the aiohttp session to free resources.
        """
        if self._session and not self._session.closed:
            await self._session.close()
            self._session = None

    @lru_cache(maxsize=100)
    def _search_model_name(self, model: str, limit: int = 5) -> tuple:
        """
        Perform regex search on model names to find all matches, cached for efficiency.

        Args:
            model (str): User-provided model name (can be partial or regex pattern).
            limit (int): Maximum number of results to return.

        Returns:
            tuple: Tuple of dictionaries with model names and their input/output cost details.

        Raises:
            KeyError: If no matching models are found.
            ValueError: If model name is empty.
        """
        if not model:
            raise ValueError("Model name cannot be empty.")

        model = model.lower().strip()
        matches = []

        if model in self.model_names:
            return ({"model": model,
                     "input_cost_per_token": self.model_prices[model]["input_cost_per_token"],
                     "output_cost_per_token": self.model_prices[model]["output_cost_per_token"]},)

        try:
            pattern = re.compile(model, re.IGNORECASE)
            matches = [name for name in self.model_names if pattern.search(name)]
        except re.error:
            matches = [name for name in self.model_names if model in name.lower()]

        matches = matches[:limit]

        if not matches:
            raise KeyError(f"No models found matching '{model}'.")

        return tuple(
            {
                "model": matched_model,
                "input_cost_per_token": self.model_prices[matched_model]["input_cost_per_token"],
                "output_cost_per_token": self.model_prices[matched_model]["output_cost_per_token"]
            }
            for matched_model in matches
        )

    def calculate_prompt_cost(self, model: str) -> Dict:
        """
        Return input/output costs per token for all matching models.

        Args:
            model (str): Model name (can be partial or regex pattern).

        Returns:
            Dict: JSON object with model names and input/output costs per token.
        """
        matched_models = self._search_model_name(model)
        return {match["model"]: {
            "input_cost_per_token": match["input_cost_per_token"],
            "output_cost_per_token": match["output_cost_per_token"]
        } for match in matched_models}

    def count_tokens(self, model: str) -> Dict:
        """
        Return input/output costs per token for all matching models.

        Args:
            model (str): Model name (can be partial or regex pattern).

        Returns:
            Dict: JSON object with model names and input/output costs per token.
        """
        matched_models = self._search_model_name(model)
        return {match["model"]: {
            "input_cost_per_token": match["input_cost_per_token"],
            "output_cost_per_token": match["output_cost_per_token"]
        } for match in matched_models}

    def get_model_costs(self, model: str) -> Dict:
        """
        Get the input/output cost details for all matching models.

        Args:
            model (str): Model name (can be partial or regex pattern).

        Returns:
            Dict: JSON object with input/output cost details for all matching models.
        """
        matched_models = self._search_model_name(model)
        return {match["model"]: {
            "input_cost_per_token": match["input_cost_per_token"],
            "output_cost_per_token": match["output_cost_per_token"]
        } for match in matched_models}