"""Utilidades para el cálculo de costos de tokens.

Define precios por modelo y funciones para estimar costos a partir del uso de tokens
(prompt, completion, cached).
"""
from typing import Dict, Any, Optional

# Precios actualizados por 1M tokens (USD) - Enero 2026
MODEL_PRICING = {
    "gpt-4.1": {
        "input": 2.00,
        "cached_input": 0.50,
        "output": 8.00,
    },
    "gpt-4.1-mini": {
        "input": 0.40,
        "cached_input": 0.10,
        "output": 1.60,
    },
    "gpt-4.1-nano": {
        "input": 0.10,
        "cached_input": 0.03,
        "output": 0.40,
    },
    "gpt-5.1": {
        "input": 1.25,
        "cached_input": 0.13,
        "output": 10.00,
    }
}

# Legacy mapping o fase por defecto (si no se especifica modelo)
PRICING = {
    "extract": MODEL_PRICING["gpt-4.1"],
    "normalize": MODEL_PRICING["gpt-4.1-mini"],
}


def compute_cost_from_usage(
    usage: Dict[str, Any] | None, 
    phase: str = "extract", 
    model: Optional[str] = None
) -> Dict[str, Any] | None:
    """Calcula el costo estimado basado en el uso de tokens, la fase o el modelo específico."""
    if not isinstance(usage, dict):
        return None
    
    prompt_tokens = int(usage.get("prompt_tokens") or 0)
    completion_tokens = int(usage.get("completion_tokens") or 0)
    
    cached_tokens = 0
    try:
        details = usage.get("prompt_tokens_details") or {}
        cached_tokens = int(details.get("cached_tokens") or 0)
    except Exception:
        cached_tokens = 0
        
    non_cached_prompt = max(prompt_tokens - cached_tokens, 0)
    
    # Determinar tarifas
    rates = None
    if model and model in MODEL_PRICING:
        rates = MODEL_PRICING[model]
    else:
        # Fallback a fase (backward compatibility)
        rates = PRICING.get(phase, PRICING["extract"])
    
    rate_in = rates["input"]
    rate_cached = rates["cached_input"]
    rate_out = rates["output"]
    
    cost_input = (non_cached_prompt / 1_000_000.0) * rate_in
    cost_cached = (cached_tokens / 1_000_000.0) * rate_cached
    cost_output = (completion_tokens / 1_000_000.0) * rate_out
    
    return {
        "input_usd": round(cost_input, 6),
        "cached_input_usd": round(cost_cached, 6),
        "output_usd": round(cost_output, 6),
        "total_usd": round(cost_input + cost_cached + cost_output, 6),
    }


def accumulate_usage(agg: Dict[str, Any], new_usage: Dict[str, Any] | None) -> None:
    """Acumula tokens de `new_usage` en el dict agregado `agg` (in-place)."""
    if not isinstance(new_usage, dict):
        return
    agg["prompt_tokens"] += int(new_usage.get("prompt_tokens") or 0)
    agg["completion_tokens"] += int(new_usage.get("completion_tokens") or 0)
    details = new_usage.get("prompt_tokens_details") or {}
    agg["prompt_tokens_details"]["cached_tokens"] += int(details.get("cached_tokens") or 0)


def sum_costs(a: Dict[str, Any] | None, b: Dict[str, Any] | None) -> Dict[str, Any]:
    """Suma dos objetos de costos."""
    def safe(x): return float(x) if x is not None else 0.0
    
    input_usd = safe((a or {}).get("input_usd")) + safe((b or {}).get("input_usd"))
    cached_usd = safe((a or {}).get("cached_input_usd")) + safe((b or {}).get("cached_input_usd"))
    output_usd = safe((a or {}).get("output_usd")) + safe((b or {}).get("output_usd"))
    
    total = {
        "input_usd": round(input_usd, 6),
        "cached_input_usd": round(cached_usd, 6),
        "output_usd": round(output_usd, 6),
    }
    total["total_usd"] = round(total["input_usd"] + total["cached_input_usd"] + total["output_usd"], 6)
    return total