#!/usr/bin/env python3
"""Backfill cost data in router_data.csv using cached responses.

This script computes costs for each (query, llm) pair by:
1. Looking up the cached response for token estimation
2. Applying model-specific output pricing

Usage:
    python -m embergraphrouter.data_processing.backfill_costs
"""

import hashlib
import json
import logging
from pathlib import Path

import pandas as pd
import yaml

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

_MODULE_DIR = Path(__file__).resolve().parent.parent

# LLM name mapping from router_data.csv names to pricing model IDs
# Format: display_name -> (provider, model_id)
LLM_NAME_TO_PRICING = {
    "Gemini 2.5 Flash Lite": ("google", "gemini-2.5-flash-lite"),
    "Claude Haiku 4.5": ("anthropic", "claude-haiku-4.5"),
    "GPT-5 Nano": ("openai", "gpt-5-nano"),
    "GPT-5.2": ("openai", "gpt-5.2"),
    "GPT-5.2 Codex": ("openai", "gpt-5.2-codex"),
    "Claude Opus 4.6": ("anthropic", "claude-opus-4.6"),
    "Gemini 3 Flash Preview": ("google", "gemini-3-flash-preview"),
    "Gemini 3 Pro Preview": ("google", "gemini-3-pro-preview"),
    "Gemini 3.1 Pro Preview": ("google", "gemini-3.1-pro-preview"),
}


def _query_hash(query: str, llm_name: str) -> str:
    """Generate a unique hash for a query+llm pair."""
    return hashlib.sha256(f"{query}||{llm_name}".encode()).hexdigest()[:16]


def load_pricing() -> dict[str, dict[str, dict]]:
    """Load pricing data from the Ember pricing YAML."""
    pricing_path = _MODULE_DIR.parent / "src" / "ember" / "models" / "pricing" / "pricing.yaml"
    if not pricing_path.exists():
        logger.warning("Pricing file not found at %s, using defaults", pricing_path)
        return {}

    with open(pricing_path, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    return data.get("providers", {})


def estimate_token_count(text: str) -> int:
    """Estimate token count using a simple heuristic.

    Approximation: ~4 characters per token for English text.
    This avoids dependency on tiktoken while being reasonably accurate.
    """
    if not text:
        return 0
    return max(1, len(text) // 4)


def compute_cost(
    response_text: str,
    llm_name: str,
    pricing: dict[str, dict[str, dict]],
) -> float:
    """Compute the cost for a single response.

    Args:
        response_text: The LLM response text
        llm_name: Display name of the LLM (e.g., "GPT-5.1")
        pricing: Loaded pricing data from YAML

    Returns:
        Cost in USD
    """
    mapping = LLM_NAME_TO_PRICING.get(llm_name)
    if not mapping:
        logger.warning("Unknown LLM name: %s", llm_name)
        return 0.0

    provider, model_id = mapping
    provider_data = pricing.get(provider, {}).get("models", {})
    model_pricing = provider_data.get(model_id)

    if not model_pricing:
        logger.warning("No pricing found for %s/%s", provider, model_id)
        return 0.0

    # Output cost per 1M tokens
    output_price_per_million = model_pricing.get("output", 0.0)

    # Estimate output tokens
    output_tokens = estimate_token_count(response_text)

    # Cost = (tokens / 1,000,000) * price_per_million
    cost = (output_tokens / 1_000_000) * output_price_per_million

    return cost


def backfill_costs(
    router_data_path: Path,
    cache_path: Path,
    output_path: Path | None = None,
) -> None:
    """Backfill cost data in router_data.csv.

    Args:
        router_data_path: Path to router_data.csv
        cache_path: Path to response_cache.json
        output_path: Path to write updated CSV (defaults to overwriting input)
    """
    if output_path is None:
        output_path = router_data_path

    # Load data
    logger.info("Loading router data from %s...", router_data_path)
    df = pd.read_csv(router_data_path)
    original_cost_sum = df["cost"].sum()
    logger.info("  Rows: %d, Current cost sum: %.4f", len(df), original_cost_sum)

    logger.info("Loading response cache from %s...", cache_path)
    with open(cache_path, "r", encoding="utf-8") as f:
        cache = json.load(f)
    logger.info("  Cache entries: %d", len(cache))

    # Load pricing
    pricing = load_pricing()
    if not pricing:
        logger.error("Failed to load pricing data")
        return

    # Log pricing info
    logger.info("Pricing loaded for providers: %s", list(pricing.keys()))

    # Compute costs
    logger.info("Computing costs...")
    costs = []
    cache_misses = 0

    for idx, row in df.iterrows():
        query = row["query"]
        llm_name = row["llm"]
        cache_key = _query_hash(query, llm_name)

        if cache_key in cache:
            response = cache[cache_key].get("response", "")
            cost = compute_cost(response, llm_name, pricing)
        else:
            cache_misses += 1
            cost = 0.0

        costs.append(cost)

        # Progress
        if (idx + 1) % 5000 == 0:
            logger.info("  Processed %d/%d rows...", idx + 1, len(df))

    df["cost"] = costs

    if cache_misses > 0:
        logger.warning("Cache misses: %d (%.1f%%)", cache_misses, cache_misses / len(df) * 100)

    # Show cost distribution before normalization
    logger.info("Cost distribution (raw, USD):")
    logger.info("  Mean: $%.6f", df["cost"].mean())
    logger.info("  Max: $%.6f", df["cost"].max())
    logger.info("  Sum: $%.4f", df["cost"].sum())

    # Show per-LLM costs
    logger.info("\nPer-LLM cost breakdown (raw):")
    for llm in df["llm"].unique():
        llm_df = df[df["llm"] == llm]
        avg_cost = llm_df["cost"].mean()
        total_cost = llm_df["cost"].sum()
        logger.info("  %s: avg=$%.6f, total=$%.4f", llm, avg_cost, total_cost)

    # Normalize costs per task (min-max scaling)
    logger.info("\nNormalizing costs per task...")

    def _minmax(series: pd.Series) -> pd.Series:
        maximum = series.max()
        minimum = series.min()
        denom = maximum - minimum
        if denom == 0:
            return pd.Series(0.0, index=series.index)
        return (series - minimum) / denom

    df["cost"] = df.groupby("task_id")["cost"].transform(_minmax)

    # Save
    df.to_csv(output_path, index=False)
    logger.info("\nSaved updated data to %s", output_path)

    # Summary
    logger.info("\n" + "=" * 60)
    logger.info("COST BACKFILL COMPLETE")
    logger.info("=" * 60)
    logger.info("Normalized cost stats:")
    logger.info("  Mean: %.4f", df["cost"].mean())
    logger.info("  Non-zero: %d (%.1f%%)", (df["cost"] > 0).sum(), (df["cost"] > 0).mean() * 100)


if __name__ == "__main__":
    config_path = _MODULE_DIR / "configs" / "config.yaml"
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    # Resolve paths
    for key in ("saved_router_data_path",):
        if key in config and not Path(config[key]).is_absolute():
            config[key] = str(_MODULE_DIR / config[key])

    cache_dir = _MODULE_DIR / "data" / ".router_cache"

    backfill_costs(
        router_data_path=Path(config["saved_router_data_path"]),
        cache_path=cache_dir / "response_cache.json",
    )
