"""
Cost Calculator

Calculates LLM usage costs based on token usage and pricing data.
Integrates with pricing fetcher and cache systems.
"""

from typing import Any, Dict, List, Optional

from loguru import logger

from .pricing_cache import PricingCache
from .pricing_fetcher import PricingFetcher


class CostCalculator:
    """Calculates LLM usage costs."""

    def __init__(self, cache_dir: Optional[str] = None):
        self.cache = PricingCache(cache_dir)
        self.pricing_fetcher = None

    async def __aenter__(self):
        self.pricing_fetcher = PricingFetcher()
        await self.pricing_fetcher.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.pricing_fetcher:
            await self.pricing_fetcher.__aexit__(exc_type, exc_val, exc_tb)

    async def get_model_pricing(
        self, model_name: str, provider: str = None
    ) -> Dict[str, float]:
        """Get pricing for a model and provider (cached or fetched)."""
        # Create cache key that includes provider
        cache_key = f"{provider}:{model_name}" if provider else model_name

        # Try cache first
        cached_pricing = self.cache.get(f"model:{cache_key}")
        if cached_pricing:
            return cached_pricing

        # Fetch from API
        if self.pricing_fetcher:
            pricing = await self.pricing_fetcher.get_model_pricing(
                model_name, provider
            )
            if pricing:
                self.cache.set(f"model:{cache_key}", pricing)
                return pricing

        # No pricing found
        logger.warning(
            f"No pricing found for {model_name} (provider: {provider})"
        )
        return None

    async def calculate_cost(
        self,
        model_name: str,
        prompt_tokens: int,
        completion_tokens: int,
        provider: str = None,
    ) -> Dict[str, float]:
        """
        Calculate cost for a single LLM call.

        Returns:
            Dict with prompt_cost, completion_cost, total_cost
        """
        pricing = await self.get_model_pricing(model_name, provider)

        # If no pricing found, return zero cost
        if pricing is None:
            return {
                "prompt_cost": 0.0,
                "completion_cost": 0.0,
                "total_cost": 0.0,
                "pricing_used": None,
                "error": "No pricing data available for this model",
            }

        # Convert tokens to thousands for pricing calculation
        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        total_cost = prompt_cost + completion_cost

        return {
            "prompt_cost": round(prompt_cost, 6),
            "completion_cost": round(completion_cost, 6),
            "total_cost": round(total_cost, 6),
            "pricing_used": pricing,
        }

    async def calculate_batch_costs(
        self, usage_records: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Calculate costs for multiple usage records.

        Expected record format:
        {
            "model_name": str,
            "provider": str (optional),
            "prompt_tokens": int,
            "completion_tokens": int,
            "research_id": int (optional),
            "timestamp": datetime (optional)
        }
        """
        results = []

        for record in usage_records:
            try:
                cost_data = await self.calculate_cost(
                    record["model_name"],
                    record["prompt_tokens"],
                    record["completion_tokens"],
                    record.get("provider"),
                )

                result = {**record, **cost_data}
                results.append(result)

            except Exception:
                logger.exception(
                    f"Failed to calculate cost for record {record}"
                )
                # Add record with zero cost on error
                results.append(
                    {
                        **record,
                        "prompt_cost": 0.0,
                        "completion_cost": 0.0,
                        "total_cost": 0.0,
                        "error": "Cost calculation failed",
                    }
                )

        return results

    def calculate_cost_sync(
        self, model_name: str, prompt_tokens: int, completion_tokens: int
    ) -> Dict[str, float]:
        """
        Synchronous cost calculation using cached pricing only.
        Fallback for when async is not available.
        """
        # Use cached pricing only
        pricing = self.cache.get_model_pricing(model_name)
        if not pricing:
            # Use static fallback with exact matching only
            fetcher = PricingFetcher()
            # Try exact match
            pricing = fetcher.static_pricing.get(model_name)
            if not pricing:
                # Try exact match without provider prefix
                if "/" in model_name:
                    model_only = model_name.split("/")[-1]
                    pricing = fetcher.static_pricing.get(model_only)

        # If no pricing found, return zero cost
        if not pricing:
            return {
                "prompt_cost": 0.0,
                "completion_cost": 0.0,
                "total_cost": 0.0,
                "pricing_used": None,
                "error": "No pricing data available for this model",
            }

        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        total_cost = prompt_cost + completion_cost

        return {
            "prompt_cost": round(prompt_cost, 6),
            "completion_cost": round(completion_cost, 6),
            "total_cost": round(total_cost, 6),
            "pricing_used": pricing,
        }

    async def get_research_cost_summary(
        self, usage_records: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Get cost summary for research session(s).
        """
        costs = await self.calculate_batch_costs(usage_records)

        total_cost = sum(c["total_cost"] for c in costs)
        total_prompt_cost = sum(c["prompt_cost"] for c in costs)
        total_completion_cost = sum(c["completion_cost"] for c in costs)

        total_prompt_tokens = sum(r["prompt_tokens"] for r in usage_records)
        total_completion_tokens = sum(
            r["completion_tokens"] for r in usage_records
        )
        total_tokens = total_prompt_tokens + total_completion_tokens

        # Model breakdown
        model_costs = {}
        for cost in costs:
            model = cost["model_name"]
            if model not in model_costs:
                model_costs[model] = {
                    "total_cost": 0.0,
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "calls": 0,
                }

            model_costs[model]["total_cost"] += cost["total_cost"]
            model_costs[model]["prompt_tokens"] += cost["prompt_tokens"]
            model_costs[model]["completion_tokens"] += cost["completion_tokens"]
            model_costs[model]["calls"] += 1

        return {
            "total_cost": round(total_cost, 6),
            "prompt_cost": round(total_prompt_cost, 6),
            "completion_cost": round(total_completion_cost, 6),
            "total_tokens": total_tokens,
            "prompt_tokens": total_prompt_tokens,
            "completion_tokens": total_completion_tokens,
            "total_calls": len(usage_records),
            "model_breakdown": model_costs,
            "avg_cost_per_call": (
                round(total_cost / len(usage_records), 6)
                if usage_records
                else 0.0
            ),
            "cost_per_token": (
                round(total_cost / total_tokens, 8) if total_tokens > 0 else 0.0
            ),
        }
