"""
AI Usage Tracker - Sledovani usage a nakladu pro vsechny AI poskytovatele.
Rozsiruje stavajici AI usage system o detailni tracking podle poskytovatelu a modelu.
"""

import json
import datetime
import os
import xbmc
import xbmcaddon
import traceback

from wcs.user_data import _load_json_data, _store_json_data, _profile

# Cesta k JSON souboru pro AI usage data
_AI_USAGE_FILE = os.path.join(_profile, 'ai_usage_stats.json')

# Cenové tabulky pro jednotlivé poskytovatele ($/1M tokenů)
PRICING_TABLES = {
    "openai": {
        "gpt-5.2": {"input": 1.75, "cached_input": 0.175, "output": 14.00},
        "gpt-5.1": {"input": 1.25, "cached_input": 0.125, "output": 10.00},
        "gpt-5-mini": {"input": 0.25, "cached_input": 0.025, "output": 2.00},
        "gpt-4.1": {"input": 2.00, "cached_input": 1.00, "output": 8.00},
        "gpt-4.1-mini": {"input": 0.40, "cached_input": 0.20, "output": 1.60},
        "gpt-4.1-nano": {"input": 0.10, "cached_input": 0.05, "output": 0.40},
        "gpt-4o": {"input": 2.50, "cached_input": 1.25, "output": 10.00},
        "gpt-4o-latest": {"input": 2.50, "cached_input": 1.25, "output": 10.00},
        "chatgpt-4o-latest": {"input": 2.50, "cached_input": 1.25, "output": 10.00},
        "gpt-4o-mini": {"input": 0.15, "cached_input": 0.075, "output": 0.30},
        "o3": {"input": 1.00, "output": 4.00},
        "o4-mini": {"input": 0.60, "output": 2.40},
        "o3-mini": {"input": 0.15, "cached_input": 0.075, "output": 0.60},
        "o1": {"input": 15.00, "cached_input": 7.50, "output": 60.00},
        "o1-mini": {"input": 3.00, "cached_input": 1.50, "output": 12.00},
        "gpt-35-turbo": {"input": 0.50, "output": 1.50}
    },
    "anthropic": {
        "claude-opus-4-20250514": {"input": 15.00, "cached_input": 0.30, "output": 75.00},
        "claude-sonnet-4-20250514": {"input": 3.00, "cached_input": 0.30, "output": 15.00},
        "claude-3-5-sonnet-20241022": {"input": 3.00, "cached_input": 0.30, "output": 15.00},
        "claude-3-5-haiku-20241022": {"input": 0.80, "cached_input": 0.08, "output": 4.00},
        "claude-3-opus-20240229": {"input": 15.00, "cached_input": 1.50, "output": 75.00},
        "claude-3-sonnet-20240229": {"input": 3.00, "cached_input": 0.30, "output": 15.00},
        "claude-3-haiku-20240307": {"input": 0.25, "cached_input": 0.025, "output": 1.25}
    },
    "google": {
        "gemini-2.5-pro": {"input": 1.25, "cached_input": 0.31, "output": 10.00},
        "gemini-2.5-flash": {"input": 0.30, "cached_input": 0.075, "output": 2.50},
        "gemini-2.5-flash-lite": {"input": 0.10, "cached_input": 0.025, "output": 0.40},
        "gemini-2.0-flash": {"input": 0.15, "cached_input": 0.0375, "output": 0.60},
        "gemini-2.0-flash-lite": {"input": 0.075, "cached_input": 0.0188, "output": 0.30},
        "gemini-1.5-pro": {"input": 1.25, "cached_input": 0.31, "output": 5.00},
        "gemini-1.5-flash": {"input": 0.075, "cached_input": 0.0188, "output": 0.30},
        "gemini-1.5-flash-8b": {"input": 0.0375, "cached_input": 0.0094, "output": 0.15}
    },
    "mistral": {
        "mistral-large-latest": {"input": 2.00, "output": 6.00},
        "magistral-medium-latest": {"input": 2.00, "output": 5.00},
        "mistral-medium-latest": {"input": 0.40, "output": 2.00},
        "magistral-small-latest": {"input": 0.10, "output": 0.30},
        "mistral-small-latest": {"input": 0.10, "output": 0.30},
        "codestral-latest": {"input": 0.30, "output": 0.90},
        "devstral-small-latest": {"input": 0.10, "output": 0.30},
        "pixtral-large-latest": {"input": 2.00, "output": 6.00},
        "ministral-8b-latest": {"input": 0.10, "output": 0.10},
        "ministral-3b-latest": {"input": 0.04, "output": 0.04},
        "open-mistral-nemo": {"input": 0.10, "output": 0.30},
        "mistral-saba-latest": {"input": 0.10, "output": 0.30}
    },
    "groq": {
        "meta-llama/llama-4-maverick-17b-128e-instruct": {"input": 0.50, "output": 0.77},
        "meta-llama/llama-4-scout-17b-16e-instruct": {"input": 0.11, "output": 0.34},
        "llama-3.3-70b-versatile": {"input": 0.59, "output": 0.79},
        "llama-3.1-8b-instant": {"input": 0.05, "output": 0.08},
        "openai/gpt-oss-120b": {"input": 0.15, "cached_input": 0.075, "output": 0.60},
        "openai/gpt-oss-20b": {"input": 0.075, "output": 0.30},
        "groq/compound": {"input": 0.30, "output": 0.90},
        "groq/compound-mini": {"input": 0.10, "output": 0.30},
        "moonshotai/kimi-k2-instruct-0905": {"input": 1.00, "cached_input": 0.50, "output": 3.00},
        "qwen/qwen3-32b": {"input": 0.29, "output": 0.59}
    }
}

def get_default_usage_structure():
    """Vrací výchozí strukturu pro sledování usage."""
    return {
        "version": "1.0",
        "created": datetime.datetime.now().isoformat(),
        "last_updated": datetime.datetime.now().isoformat(),
        "providers": {
            "openai": {"total_requests": 0, "total_cost": 0.0, "models": {}},
            "anthropic": {"total_requests": 0, "total_cost": 0.0, "models": {}},
            "google": {"total_requests": 0, "total_cost": 0.0, "models": {}},
            "mistral": {"total_requests": 0, "total_cost": 0.0, "models": {}},
            "groq": {"total_requests": 0, "total_cost": 0.0, "models": {}}
        },
        "monthly_stats": {},
        "daily_stats": {},
        "call_history": []
    }

def _migrate_usage_from_settings(addon):
    """Migruje usage data ze settings do JSON (jednorazove)."""
    try:
        settings_data = addon.getSetting('ai_usage_data')
        if settings_data:
            data = json.loads(settings_data)
            if isinstance(data, dict) and "providers" in data:
                _store_json_data(_AI_USAGE_FILE, data)
                addon.setSetting('ai_usage_data', '')
                xbmc.log('[wcs] Migrated ai_usage_data from settings to ai_usage_stats.json', xbmc.LOGINFO)
                return data
    except Exception as e:
        xbmc.log(f'[wcs] Error migrating usage data: {e}', xbmc.LOGERROR)
        traceback.print_exc()
    return None

def load_usage_data():
    """Nacte data o usage z JSON souboru s automatickou migraci ze settings."""
    addon = xbmcaddon.Addon()
    
    # Migrace ze settings (pouze pokud JSON neexistuje)
    if not os.path.exists(_AI_USAGE_FILE):
        migrated = _migrate_usage_from_settings(addon)
        if migrated:
            return migrated
    
    # Pouzit genericke nacteni s default hodnotou
    data = _load_json_data(_AI_USAGE_FILE, None, addon, default_value=get_default_usage_structure())
    
    # Validace struktury
    if isinstance(data, dict) and "providers" in data:
        return data
    return get_default_usage_structure()

def save_usage_data(data):
    """Ulozi data o usage do JSON souboru."""
    try:
        data["last_updated"] = datetime.datetime.now().isoformat()
        _store_json_data(_AI_USAGE_FILE, data)
        return True
    except Exception as e:
        xbmc.log(f'[wcs] Error saving ai_usage_stats.json: {e}', xbmc.LOGERROR)
        traceback.print_exc()
        return False

def calculate_cost(provider, model, usage_data):
    """
    Vypočítá náklady pro daný provider/model/usage.
    
    Args:
        provider (str): AI poskytovatel (openai, anthropic, google, mistral, groq)
        model (str): Název modelu
        usage_data (dict): Usage data s tokeny
        
    Returns:
        dict: {"input_cost": 0.0, "cached_input_cost": 0.0, "output_cost": 0.0, "total_cost": 0.0}
    """
    # Extrakce token counts z usage_data - podpora nového i starého formátu
    input_tokens = 0
    output_tokens = 0
    cached_tokens = 0

    # Nový OpenAI formát
    if "input_tokens" in usage_data:
        input_tokens = usage_data.get("input_tokens", 0)
        output_tokens = usage_data.get("output_tokens", 0)
        
        # Cached tokeny z input_tokens_details
        if "input_tokens_details" in usage_data:
            cached_tokens = usage_data["input_tokens_details"].get("cached_tokens", 0)
            input_tokens = max(0, input_tokens - cached_tokens) # Odečteme cached tokeny z normálních
    
    # Starý formát (fallback)
    else:
        input_tokens = usage_data.get("prompt_tokens", 0)
        output_tokens = usage_data.get("completion_tokens", 0)
        
        # Speciální handling pro cached tokeny (starý OpenAI formát)
        if "prompt_tokens_details" in usage_data:
            cached_tokens = usage_data["prompt_tokens_details"].get("cached_tokens", 0)
            input_tokens = max(0, input_tokens - cached_tokens) # Odečteme cached tokeny z normálních

    # Získání ceníku
    pricing = PRICING_TABLES.get(provider, {}).get(model, {})
    
    if not pricing:
        xbmc.log(f'[wcs] No pricing data for {provider}/{model}, returning 0 cost but tracked tokens', xbmc.LOGWARNING)
        return {
            "input_cost": 0.0, 
            "cached_input_cost": 0.0, 
            "output_cost": 0.0, 
            "total_cost": 0.0, 
            "input_tokens": input_tokens, 
            "cached_tokens": cached_tokens, 
            "output_tokens": output_tokens
        }
    
    # Výpočet nákladů (ceny jsou za 1M tokenů)
    input_cost = (input_tokens / 1_000_000) * pricing.get("input", 0.0)
    cached_input_cost = (cached_tokens / 1_000_000) * pricing.get("cached_input", 0.0)
    output_cost = (output_tokens / 1_000_000) * pricing.get("output", 0.0)
    total_cost = input_cost + cached_input_cost + output_cost
    
    # Debug log výsledku
    xbmc.log(f'[wcs] Cost calculation: input={input_tokens}, cached={cached_tokens}, output={output_tokens}, total_cost=${total_cost:.6f}', xbmc.LOGINFO)
    
    return {
        "input_cost": input_cost,
        "cached_input_cost": cached_input_cost,
        "output_cost": output_cost,
        "total_cost": total_cost,
        "input_tokens": input_tokens,
        "cached_tokens": cached_tokens,
        "output_tokens": output_tokens
    }

def track_ai_usage(provider, model, usage_data, request_type="standard", session_id=None):
    """
    Zaznamenává použití AI API s detailním trackingem nákladů.
    
    Args:
        provider (str): AI poskytovatel
        model (str): Použitý model
        usage_data (dict): Usage data z API response
        request_type (str): Typ požadavku (standard, game, recommendation, etc.)
        session_id (str): Volitelné ID herní session. Pokud je zadáno,
            kumuluje se do jednoho záznamu v call_history místo vytváření nových.
    """
    try:
        data = load_usage_data()
        now = datetime.datetime.now()
        today = now.strftime("%Y-%m-%d")
        month = now.strftime("%Y-%m")
        
        # Výpočet nákladů
        cost_info = calculate_cost(provider, model, usage_data)
        
        # Update provider stats
        if provider not in data["providers"]:
            data["providers"][provider] = {"total_requests": 0, "total_cost": 0.0, "models": {}}
        
        provider_data = data["providers"][provider]
        provider_data["total_requests"] += 1
        provider_data["total_cost"] += cost_info["total_cost"]
        
        # Update model stats
        if model not in provider_data["models"]:
            provider_data["models"][model] = {
                "requests": 0,
                "total_cost": 0.0,
                "input_tokens": 0,
                "cached_tokens": 0,
                "output_tokens": 0,
                "request_types": {}
            }
        
        model_data = provider_data["models"][model]
        model_data["requests"] += 1
        model_data["total_cost"] += cost_info["total_cost"]
        model_data["input_tokens"] += cost_info["input_tokens"]
        model_data["cached_tokens"] += cost_info["cached_tokens"]
        model_data["output_tokens"] += cost_info["output_tokens"]
        
        # Track request types
        if request_type not in model_data["request_types"]:
            model_data["request_types"][request_type] = {"count": 0, "cost": 0.0}
        
        model_data["request_types"][request_type]["count"] += 1
        model_data["request_types"][request_type]["cost"] += cost_info["total_cost"]
        
        # Monthly stats
        if month not in data["monthly_stats"]:
            data["monthly_stats"][month] = {"requests": 0, "cost": 0.0, "providers": {}}
        
        monthly_data = data["monthly_stats"][month]
        monthly_data["requests"] += 1
        monthly_data["cost"] += cost_info["total_cost"]
        
        if provider not in monthly_data["providers"]:
            monthly_data["providers"][provider] = {"requests": 0, "cost": 0.0}
        
        monthly_data["providers"][provider]["requests"] += 1
        monthly_data["providers"][provider]["cost"] += cost_info["total_cost"]
        
        # Daily stats
        if today not in data["daily_stats"]:
            data["daily_stats"][today] = {"requests": 0, "cost": 0.0, "providers": {}}
        
        daily_data = data["daily_stats"][today]
        daily_data["requests"] += 1
        daily_data["cost"] += cost_info["total_cost"]
        
        if provider not in daily_data["providers"]:
            daily_data["providers"][provider] = {"requests": 0, "cost": 0.0}
        
        daily_data["providers"][provider]["requests"] += 1
        daily_data["providers"][provider]["cost"] += cost_info["total_cost"]
        
        # Inicializace call_history pokud neexistuje (migrace ze starých dat)
        if "call_history" not in data:
            data["call_history"] = []
        
        # Session-based kumulace: pokud je session_id, akumuluj do existujícího
        if session_id:
            existing = None
            for record in data["call_history"]:
                if record.get("session_id") == session_id:
                    existing = record
                    break
            
            if existing:
                # Akumulace do existujícího záznamu
                existing["cost"] += cost_info["total_cost"]
                existing["input_tokens"] += cost_info["input_tokens"]
                existing["cached_tokens"] += cost_info["cached_tokens"]
                existing["output_tokens"] += cost_info["output_tokens"]
                existing["total_tokens"] += (
                    cost_info["input_tokens"] + cost_info["cached_tokens"] + cost_info["output_tokens"]
                )
                existing["inference_count"] = existing.get("inference_count", 1) + 1
                existing["last_updated"] = now.isoformat()
            else:
                # Nový session záznam
                call_record = {
                    "timestamp": now.isoformat(),
                    "last_updated": now.isoformat(),
                    "provider": provider,
                    "model": model,
                    "request_type": request_type,
                    "session_id": session_id,
                    "inference_count": 1,
                    "cost": cost_info["total_cost"],
                    "input_tokens": cost_info["input_tokens"],
                    "cached_tokens": cost_info["cached_tokens"],
                    "output_tokens": cost_info["output_tokens"],
                    "total_tokens": cost_info["input_tokens"] + cost_info["cached_tokens"] + cost_info["output_tokens"]
                }
                data["call_history"].append(call_record)
        else:
            # Standardní individuální záznam (bez session)
            call_record = {
                "timestamp": now.isoformat(),
                "provider": provider,
                "model": model,
                "request_type": request_type,
                "cost": cost_info["total_cost"],
                "input_tokens": cost_info["input_tokens"],
                "cached_tokens": cost_info["cached_tokens"],
                "output_tokens": cost_info["output_tokens"],
                "total_tokens": cost_info["input_tokens"] + cost_info["cached_tokens"] + cost_info["output_tokens"]
            }
            data["call_history"].append(call_record)
        
        # Omezíme historii na posledních 500 záznamů (aby soubor nebyl příliš velký)
        if len(data["call_history"]) > 500:
            data["call_history"] = data["call_history"][-500:]
        
        # Uložíme data
        success = save_usage_data(data)
        
        if success:
            xbmc.log(f'[wcs] AI usage tracked: {provider}/{model} - ${cost_info["total_cost"]:.6f}', xbmc.LOGINFO)
        
        return success
        
    except Exception as e:
        xbmc.log(f'[wcs] Error tracking AI usage: {e}', xbmc.LOGERROR)
        traceback.print_exc()
        return False

def get_usage_summary():
    """
    Vrací shrnutí usage a nákladů.
    
    Returns:
        dict: Detailní shrnutí s celkovými náklady, breakdown podle provideru atd.
    """
    try:
        data = load_usage_data()
        
        # Celkové statistiky
        total_requests = sum(p["total_requests"] for p in data["providers"].values())
        total_cost = sum(p["total_cost"] for p in data["providers"].values())
        
        # Dnešní statistiky
        today = datetime.datetime.now().strftime("%Y-%m-%d")
        today_stats = data["daily_stats"].get(today, {"requests": 0, "cost": 0.0, "providers": {}})
        
        # Měsíční statistiky
        month = datetime.datetime.now().strftime("%Y-%m")
        month_stats = data["monthly_stats"].get(month, {"requests": 0, "cost": 0.0, "providers": {}})
        
        # Nejpoužívanější modely s token daty
        all_models = []
        for provider, provider_data in data["providers"].items():
            for model, model_data in provider_data["models"].items():
                all_models.append({
                    "provider": provider,
                    "model": model,
                    "requests": model_data["requests"],
                    "cost": model_data["total_cost"],
                    "avg_cost_per_request": model_data["total_cost"] / max(model_data["requests"], 1),
                    "input_tokens": model_data.get("input_tokens", 0),
                    "cached_tokens": model_data.get("cached_tokens", 0),
                    "output_tokens": model_data.get("output_tokens", 0)
                })
        
        # Seřazení podle nákladů
        top_models = sorted(all_models, key=lambda x: x["cost"], reverse=True)[:10]
        
        summary = {
            "totals": {
                "requests": total_requests,
                "cost": total_cost,
                "avg_cost_per_request": total_cost / max(total_requests, 1)
            },
            "today": today_stats,
            "this_month": month_stats,
            "providers": data["providers"],
            "top_models": top_models,
            "last_updated": data.get("last_updated", "")
        }
        
        return summary
        
    except Exception as e:
        xbmc.log(f'[wcs] Error getting usage summary: {e}', xbmc.LOGERROR)
        traceback.print_exc()
        return {"error": str(e)}

def reset_usage_data(period="all"):
    """
    Resetuje usage data.
    
    Args:
        period (str): "all", "monthly", "daily", nebo "provider:openai" pro specifický provider
    """
    try:
        if period == "all":
            # Kompletní reset - smaže JSON soubor
            if os.path.exists(_AI_USAGE_FILE):
                os.remove(_AI_USAGE_FILE)
        else:
            # Částečný reset - načteme data, upravíme a uložíme
            data = load_usage_data()
            
            if period == "monthly":
                data["monthly_stats"] = {}
            elif period == "daily":
                data["daily_stats"] = {}
            elif period.startswith("provider:"):
                provider = period.split(":", 1)[1]
                if provider in data["providers"]:
                    data["providers"][provider] = {"total_requests": 0, "total_cost": 0.0, "models": {}}
            
            success = save_usage_data(data)
            if not success:
                return False
        
        xbmc.log(f'[wcs] AI usage data reset: {period}', xbmc.LOGINFO)
        return True
        
    except Exception as e:
        xbmc.log(f'[wcs] Error resetting usage data: {e}', xbmc.LOGERROR)
        traceback.print_exc()
        return False

def format_cost(cost, addon=None):
    """Formátuje náklady pro zobrazení podle uživatelského nastavení."""
    # Kurz koruny k dolaru
    USD_TO_CZK = 22.0
    czk_cost = cost * USD_TO_CZK
    
    # Výchozí nastavení pokud addon není k dispozici
    currency_setting = "USD(CZK)"
    if addon:
        try:
            currency_setting = addon.getSetting("ai_cost_currency") or "USD(CZK)"
            # Bezpečnostní převod indexu na text
            if currency_setting == "0":
                currency_setting = "USD"
            elif currency_setting == "1":
                currency_setting = "CZK"
            elif currency_setting == "2":
                currency_setting = "USD(CZK)"
        except:
            currency_setting = "USD(CZK)"
    
    # Formátování podle nastavení
    if currency_setting == "USD":
        if cost >= 1.0:
            return f"${cost:.2f}"
        elif cost >= 0.01:
            return f"${cost:.3f}"
        else:
            return f"${cost:.6f}"
    elif currency_setting == "CZK":
        if czk_cost >= 10.0:
            return f"{czk_cost:.0f} Kč"
        elif czk_cost >= 0.1:
            return f"{czk_cost:.1f} Kč"
        else:
            return f"{czk_cost:.2f} Kč"
    else:  # "USD(CZK)" nebo "Obojí" pro zpětnou kompatibilitu
        if cost >= 1.0:
            return f"${cost:.2f} ({czk_cost:.0f} Kč)"
        elif cost >= 0.01:
            return f"${cost:.3f} ({czk_cost:.1f} Kč)"
        else:
            return f"${cost:.6f} ({czk_cost:.2f} Kč)"

def format_tokens(tokens):
    """Formátuje počet tokenů pro zobrazení."""
    if tokens >= 1_000_000:
        return f"{tokens/1_000_000:.1f}M"
    elif tokens >= 1_000:
        return f"{tokens/1_000:.1f}K"
    else:
        return str(tokens)

def get_call_history(limit=50):
    """
    Vrací historii posledních AI volání.
    
    Args:
        limit (int): Maximální počet záznamů k vrácení
        
    Returns:
        list: Seznam individuálních volání seřazených od nejnovějších
    """
    try:
        data = load_usage_data()
        call_history = data.get("call_history", [])
        
        # Seřazení podle timestamp (nejnovější první) a omezení
        sorted_calls = sorted(call_history, key=lambda x: x.get("timestamp", ""), reverse=True)
        return sorted_calls[:limit]
        
    except Exception as e:
        xbmc.log(f'[wcs] Error getting call history: {e}', xbmc.LOGERROR)
        traceback.print_exc()
        return []
 