"""
cost_tracker.py - Track and enforce API usage budgets across all providers.

Smart routing based on remaining budgets:
- "thinking" tasks → OpenRouter (Nex-N2-Pro, free)
- "translation" tasks → Groq (Llama 3.3 70B, fast + free)
- "transcription" tasks → Groq (Whisper, free)
- "fallback" → whichever has budget

Persisted via StateManager. Alert levels: OK, WARN, CRITICAL, EXHAUSTED.
"""

import json
import os
import threading
import time
from datetime import datetime
from typing import Any, Dict, List, Optional


# ================================================================
# Provider Budget Configs
# ================================================================

PROVIDER_CONFIGS = {
    "openrouter": {
        "daily_limit_per_key": 45,  # Free tier limit per key
        "monthly_cost_free": 0.0,
        "cost_per_1k_tokens": 0.0,  # Free
        "models": ["nex-agi/nex-n2-pro:free"],
    },
    "scaleway": {
        "total_free_tokens": 1_000_000,  # 1M tokens free
        "monthly_cost_free": 0.0,
        "cost_per_1k_tokens": 0.0,  # Free tier
        "models": ["qwen3.5-397b-a17b"],
    },
    "groq": {
        "daily_limit": 14400,  # requests per day
        "monthly_cost_free": 0.0,
        "cost_per_1k_tokens": 0.0,  # Free
        "models": ["llama-3.3-70b-versatile", "whisper-large-v3"],
    },
    "siliconflow": {
        "daily_limit": 1000,  # estimated
        "cost_per_1k_chars_tts": {
            "fish-speech": 0.015,      # $15/M UTF-8 bytes
            "cosyvoice": 0.00715,      # $7.15/M bytes
            "indextts": 0.00715,       # $7.15/M bytes
        },
        "models": ["fishaudio/fish-speech-1.5", "FunAudioLLM/CosyVoice2-0.5B", "IndexTeam/IndexTTS-2"],
    },
}

ALERT_THRESHOLDS = {
    "ok": 0.0,
    "warn": 0.50,      # 50% used
    "critical": 0.80,   # 80% used
    "exhausted": 1.0,   # 100% used
}


class CostTracker:
    """Track and enforce API usage budgets across all providers."""

    SAVE_INTERVAL = 10  # Auto-save every N requests

    def __init__(self, state=None):
        self.state = state
        self._lock = threading.Lock()
        self._request_counter = 0  # For auto-save interval
        self._ensure_budget_state()

    # ================================================================
    # Public API
    # ================================================================

    def record_request(self, provider: str, model: str = "", tokens_used: int = 0,
                       key_index: Optional[int] = None, chars_used: int = 0) -> Dict[str, Any]:
        """Record an API call for budget tracking.

        Args:
            provider: "openrouter", "scaleway", "groq", "siliconflow"
            model: The model name used
            tokens_used: Approximate tokens consumed
            key_index: For OpenRouter, which key was used (0-4)
            chars_used: For TTS, characters processed

        Returns:
            Updated budget info for this provider
        """
        with self._lock:
            budget = self._get_budget()
            today = datetime.now().strftime("%Y-%m-%d")

            if provider == "openrouter":
                key_name = f"key_{key_index}" if key_index is not None else "key_0"
                if "openrouter" not in budget:
                    budget["openrouter"] = {"daily": {}, "total_requests": 0, "total_tokens": 0, "estimated_cost": 0.0}
                if today not in budget["openrouter"]["daily"]:
                    budget["openrouter"]["daily"][today] = {}
                if key_name not in budget["openrouter"]["daily"][today]:
                    budget["openrouter"]["daily"][today][key_name] = {"requests": 0, "tokens": 0}
                budget["openrouter"]["daily"][today][key_name]["requests"] += 1
                budget["openrouter"]["daily"][today][key_name]["tokens"] += tokens_used
                budget["openrouter"]["total_requests"] += 1
                budget["openrouter"]["total_tokens"] += tokens_used

            elif provider == "groq":
                if "groq" not in budget:
                    budget["groq"] = {"daily": {}, "total_requests": 0, "total_tokens": 0, "estimated_cost": 0.0}
                if today not in budget["groq"]["daily"]:
                    budget["groq"]["daily"][today] = {"requests": 0, "tokens": 0}
                budget["groq"]["daily"][today]["requests"] += 1
                budget["groq"]["daily"][today]["tokens"] += tokens_used
                budget["groq"]["total_requests"] += 1
                budget["groq"]["total_tokens"] += tokens_used

            elif provider == "scaleway":
                if "scaleway" not in budget:
                    budget["scaleway"] = {"total_tokens": 0, "estimated_cost": 0.0}
                budget["scaleway"]["total_tokens"] += tokens_used

            elif provider == "siliconflow":
                if "siliconflow" not in budget:
                    budget["siliconflow"] = {"daily": {}, "total_chars": 0, "total_cost": 0.0}
                if today not in budget["siliconflow"]["daily"]:
                    budget["siliconflow"]["daily"][today] = {"requests": 0, "chars": 0, "cost": 0.0}
                budget["siliconflow"]["daily"][today]["requests"] += 1
                budget["siliconflow"]["daily"][today]["chars"] += chars_used
                # Estimate cost
                cost_per_1k = 0.015  # Default fish-speech
                for model_name, cost in PROVIDER_CONFIGS["siliconflow"]["cost_per_1k_chars_tts"].items():
                    if model_name in model:
                        cost_per_1k = cost
                        break
                estimated_cost = (chars_used / 1000) * cost_per_1k
                budget["siliconflow"]["daily"][today]["cost"] += estimated_cost
                budget["siliconflow"]["total_chars"] += chars_used
                budget["siliconflow"]["total_cost"] += estimated_cost

            self._save_budget(budget)
            self._request_counter += 1

            # Auto-save periodically
            if self._request_counter % self.SAVE_INTERVAL == 0:
                self._persist_budget(budget)

            return self.get_remaining(provider)

    def can_make_request(self, provider: str, model: str = "") -> bool:
        """Check if the budget allows making a request to this provider."""
        remaining = self.get_remaining(provider)
        if not remaining:
            return True  # No tracking data = assume OK

        if provider == "openrouter":
            # Check if any key has remaining requests
            daily = remaining.get("daily", {})
            today = datetime.now().strftime("%Y-%m-%d")
            today_data = daily.get(today, {})
            for key_name, key_data in today_data.items():
                if key_data.get("remaining", 0) > 0:
                    return True
            return False

        elif provider == "groq":
            daily = remaining.get("daily", {})
            today = datetime.now().strftime("%Y-%m-%d")
            today_data = daily.get(today, {})
            used = today_data.get("requests", 0)
            limit = PROVIDER_CONFIGS["groq"]["daily_limit"]
            return used < limit

        elif provider == "scaleway":
            total_tokens = remaining.get("total_tokens", 0)
            limit = PROVIDER_CONFIGS["scaleway"]["total_free_tokens"]
            return total_tokens < limit

        elif provider == "siliconflow":
            return True  # Pay-per-use, always available

        return True

    def get_remaining(self, provider: str) -> Dict[str, Any]:
        """Get remaining budget for a provider."""
        budget = self._get_budget()
        today = datetime.now().strftime("%Y-%m-%d")

        if provider == "openrouter":
            or_data = budget.get("openrouter", {"daily": {}})
            daily = or_data.get("daily", {})
            today_data = daily.get(today, {})
            n_keys = len([k.strip() for k in os.getenv("OPENROUTER_API_KEYS", "").split(",") if k.strip()])
            limit_per_key = PROVIDER_CONFIGS["openrouter"]["daily_limit_per_key"]

            total_remaining = 0
            key_status = {}
            for i in range(n_keys):
                key_name = f"key_{i}"
                used = today_data.get(key_name, {}).get("requests", 0)
                remaining = max(0, limit_per_key - used)
                total_remaining += remaining
                key_status[key_name] = {
                    "used": used,
                    "limit": limit_per_key,
                    "remaining": remaining,
                }

            return {
                "provider": "openrouter",
                "total_remaining_today": total_remaining,
                "total_limit_today": n_keys * limit_per_key,
                "daily": daily,
                "key_status": key_status,
            }

        elif provider == "groq":
            groq_data = budget.get("groq", {"daily": {}})
            daily = groq_data.get("daily", {})
            today_data = daily.get(today, {})
            used = today_data.get("requests", 0)
            limit = PROVIDER_CONFIGS["groq"]["daily_limit"]
            return {
                "provider": "groq",
                "used_today": used,
                "limit_daily": limit,
                "remaining_today": max(0, limit - used),
                "pct_used": round((used / limit) * 100, 1) if limit > 0 else 0,
            }

        elif provider == "scaleway":
            sc_data = budget.get("scaleway", {})
            total_tokens = sc_data.get("total_tokens", 0)
            limit = PROVIDER_CONFIGS["scaleway"]["total_free_tokens"]
            return {
                "provider": "scaleway",
                "total_tokens_used": total_tokens,
                "total_limit": limit,
                "remaining_tokens": max(0, limit - total_tokens),
                "pct_used": round((total_tokens / limit) * 100, 1) if limit > 0 else 0,
            }

        elif provider == "siliconflow":
            sf_data = budget.get("siliconflow", {"total_cost": 0.0})
            return {
                "provider": "siliconflow",
                "total_cost": sf_data.get("total_cost", 0.0),
                "pay_per_use": True,
            }

        return {"provider": provider, "status": "unknown"}

    def get_all_budgets(self) -> Dict[str, Any]:
        """Get budget status for all providers."""
        return {
            "openrouter": self.get_remaining("openrouter"),
            "groq": self.get_remaining("groq"),
            "scaleway": self.get_remaining("scaleway"),
            "siliconflow": self.get_remaining("siliconflow"),
        }

    def get_optimal_provider(self, task_type: str) -> str:
        """Smart routing: select the best provider for a task type.

        Task types:
        - "thinking": Complex reasoning, decision-making → OpenRouter (Nex-N2-Pro)
        - "translation": EN→ES translation → Groq (Llama 3.3 70B, fast + free)
        - "transcription": Audio transcription → Groq (Whisper)
        - "tts": Text-to-speech → Supertonic (local) or SiliconFlow
        - "fallback": Any available provider
        """
        if task_type == "thinking":
            if self.can_make_request("openrouter"):
                return "openrouter"
            if self.can_make_request("scaleway"):
                return "scaleway"
            return "groq"  # Fallback for thinking

        elif task_type in ("translation", "transcription"):
            if self.can_make_request("groq"):
                return "groq"
            if self.can_make_request("openrouter"):
                return "openrouter"
            return "scaleway"

        elif task_type == "tts":
            return "supertonic"  # Local, always available

        else:  # fallback
            for provider in ["openrouter", "groq", "scaleway"]:
                if self.can_make_request(provider):
                    return provider
            return "groq"  # Last resort

    def is_budget_exhausted(self) -> bool:
        """True if ALL providers are out of budget."""
        for provider in ["openrouter", "groq", "scaleway"]:
            if self.can_make_request(provider):
                return False
        return True

    def get_daily_summary(self) -> Dict[str, Any]:
        """Get today's usage summary across all providers."""
        today = datetime.now().strftime("%Y-%m-%d")
        budget = self._get_budget()

        summary = {
            "date": today,
            "openrouter": {},
            "groq": {},
            "siliconflow": {},
            "total_estimated_cost": 0.0,
        }

        # OpenRouter
        or_daily = budget.get("openrouter", {}).get("daily", {}).get(today, {})
        or_total_reqs = sum(k.get("requests", 0) for k in or_daily.values())
        or_total_tokens = sum(k.get("tokens", 0) for k in or_daily.values())
        summary["openrouter"] = {
            "requests": or_total_reqs,
            "tokens": or_total_tokens,
            "estimated_cost": 0.0,  # Free
        }

        # Groq
        groq_daily = budget.get("groq", {}).get("daily", {}).get(today, {})
        summary["groq"] = {
            "requests": groq_daily.get("requests", 0),
            "tokens": groq_daily.get("tokens", 0),
            "estimated_cost": 0.0,  # Free
        }

        # SiliconFlow
        sf_daily = budget.get("siliconflow", {}).get("daily", {}).get(today, {})
        summary["siliconflow"] = {
            "requests": sf_daily.get("requests", 0),
            "chars": sf_daily.get("chars", 0),
            "estimated_cost": sf_daily.get("cost", 0.0),
        }
        summary["total_estimated_cost"] = sf_daily.get("cost", 0.0)

        return summary

    def reset_daily(self):
        """Reset daily counters. Called at midnight."""
        # Daily counters are keyed by date, so they auto-reset
        # Just clean up old dates
        budget = self._get_budget()
        today = datetime.now().strftime("%Y-%m-%d")

        for provider in ["openrouter", "groq", "siliconflow"]:
            if provider in budget and "daily" in budget[provider]:
                daily = budget[provider]["daily"]
                # Remove entries older than 7 days
                to_remove = [d for d in daily if d < today]
                for d in to_remove:
                    del daily[d]

        self._persist_budget(budget)
        print("[COST] Daily counters cleaned up")

    def get_alert_level(self, provider: str) -> str:
        """Get alert level for a provider: ok, warn, critical, exhausted."""
        if provider == "openrouter":
            remaining = self.get_remaining("openrouter")
            total_limit = remaining.get("total_limit_today", 1)
            total_remaining = remaining.get("total_remaining_today", 0)
            if total_limit == 0:
                return "exhausted"
            pct_used = 1.0 - (total_remaining / total_limit)
        elif provider == "groq":
            remaining = self.get_remaining("groq")
            pct_used = remaining.get("pct_used", 0) / 100.0
        elif provider == "scaleway":
            remaining = self.get_remaining("scaleway")
            pct_used = remaining.get("pct_used", 0) / 100.0
        else:
            return "ok"

        if pct_used >= 1.0:
            return "exhausted"
        elif pct_used >= 0.80:
            return "critical"
        elif pct_used >= 0.50:
            return "warn"
        return "ok"

    # ================================================================
    # Internal Methods
    # ================================================================

    def _ensure_budget_state(self):
        """Ensure provider_budget key exists in state."""
        if self.state and hasattr(self.state, "_state"):
            if "provider_budget" not in self.state._state:
                self.state._state["provider_budget"] = {}

    def _get_budget(self) -> Dict[str, Any]:
        """Get budget data from state."""
        if self.state and hasattr(self.state, "_state"):
            return self.state._state.get("provider_budget", {})
        return {}

    def _save_budget(self, budget: Dict[str, Any]):
        """Save budget data to state (in-memory only)."""
        if self.state and hasattr(self.state, "_state"):
            self.state._state["provider_budget"] = budget

    def _persist_budget(self, budget: Dict[str, Any]):
        """Persist budget data to HF Dataset."""
        self._save_budget(budget)
        if self.state:
            try:
                self.state.save()
            except Exception as e:
                print(f"[COST] Warning: could not persist budget: {e}")