#!/usr/bin/env python3
"""Regenerate router_data.csv from cached responses.

This script rebuilds router_data.csv using:
1. Queries from unified_qa_data.csv
2. Cached responses from response_cache.json
3. Correct evaluation metrics (bbeh for BBEH tasks, etc.)

No API calls are made - only cached data is used.

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

import hashlib
import json
import logging
from pathlib import Path

import pandas as pd
from tqdm import tqdm

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

# Import evaluation and embedding functions
try:
    from .llm_engine import LLMEngine
    from .utils import get_embedding, loadjson
except ImportError:
    from llm_engine import LLMEngine
    from utils import get_embedding, loadjson


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


def regenerate_router_data(
    qa_path: Path | None = None,
    cache_path: Path | None = None,
    llm_path: Path | None = None,
    output_path: Path | None = None,
    skip_missing: bool = True,
) -> pd.DataFrame:
    """Regenerate router_data.csv from cached responses.

    Args:
        qa_path: Path to unified_qa_data.csv
        cache_path: Path to response_cache.json
        llm_path: Path to LLM_Descriptions.json
        output_path: Path to write router_data.csv
        skip_missing: If True, skip queries missing from cache

    Returns:
        DataFrame with regenerated router data
    """
    # Default paths
    data_dir = _MODULE_DIR / "data"
    config_dir = _MODULE_DIR / "configs"

    qa_path = qa_path or data_dir / "unified_qa_data.csv"
    cache_path = cache_path or data_dir / ".router_cache" / "response_cache.json"
    llm_path = llm_path or config_dir / "LLM_Descriptions.json"
    output_path = output_path or data_dir / "router_data.csv"

    # Load data
    logger.info("Loading data sources...")
    qa_df = pd.read_csv(qa_path)
    llm_descriptions = loadjson(str(llm_path))
    llm_names = list(llm_descriptions.keys())

    with open(cache_path, "r", encoding="utf-8") as f:
        cache = json.load(f)

    logger.info(f"  QA queries: {len(qa_df):,}")
    logger.info(f"  LLMs: {len(llm_names)}")
    logger.info(f"  Cache entries: {len(cache):,}")

    # Initialize evaluation engine
    engine = LLMEngine(llm_names, llm_descriptions)

    # Get LLM costs from descriptions
    llm_costs = {}
    for name, desc in llm_descriptions.items():
        # Normalize cost to 0-1 range (assumes max cost ~$0.10/query)
        raw_cost = desc.get("cost", 0.01)
        llm_costs[name] = min(raw_cost / 0.10, 1.0)

    # Load existing embeddings from current router_data if available
    logger.info("Loading cached embeddings...")
    existing_router_path = data_dir / "router_data.csv"
    query_embedding_cache = {}
    task_desc_embedding_cache = {}

    if existing_router_path.exists():
        existing_df = pd.read_csv(existing_router_path)
        # Cache query embeddings
        for _, row in existing_df.drop_duplicates(subset=["query"]).iterrows():
            query_embedding_cache[row["query"]] = row["query_embedding"]
        # Cache task description embeddings
        for _, row in existing_df.drop_duplicates(subset=["task_id"]).iterrows():
            task_desc_embedding_cache[row["task_id"]] = row["task_description_embedding"]
        logger.info(f"  Loaded {len(query_embedding_cache):,} query embeddings")
        logger.info(f"  Loaded {len(task_desc_embedding_cache):,} task embeddings")

    # Compute missing task description embeddings (batch)
    logger.info("Computing task description embeddings...")
    task_desc_embeddings = {}
    missing_task_ids = []
    missing_task_descs = []

    for task_id in qa_df["task_id"].unique():
        if task_id in task_desc_embedding_cache:
            task_desc_embeddings[task_id] = task_desc_embedding_cache[task_id]
        else:
            task_desc = qa_df[qa_df["task_id"] == task_id]["task_description"].iloc[0]
            missing_task_ids.append(task_id)
            missing_task_descs.append(task_desc)

    if missing_task_descs:
        logger.info(f"  Computing {len(missing_task_descs)} task embeddings in batch...")
        embeddings = get_embedding(missing_task_descs)
        for task_id, emb in zip(missing_task_ids, embeddings):
            task_desc_embeddings[task_id] = str(emb.tolist())
        logger.info(f"  Done computing task embeddings")

    # Compute missing query embeddings (batch)
    logger.info("Computing query embeddings...")
    all_queries = qa_df["query"].unique().tolist()
    missing_queries = [q for q in all_queries if q not in query_embedding_cache]

    if missing_queries:
        logger.info(f"  Computing {len(missing_queries):,} query embeddings in batches...")
        batch_size = 100  # Process in batches to avoid timeouts
        for i in range(0, len(missing_queries), batch_size):
            batch = missing_queries[i : i + batch_size]
            embeddings = get_embedding(batch)
            for query, emb in zip(batch, embeddings):
                query_embedding_cache[query] = str(emb.tolist())
            if (i + batch_size) % 1000 == 0 or i + batch_size >= len(missing_queries):
                logger.info(f"    Progress: {min(i + batch_size, len(missing_queries)):,}/{len(missing_queries):,}")
        logger.info(f"  Done computing query embeddings")

    # Process queries
    logger.info("Regenerating router data...")
    records = []
    missing_count = 0
    error_count = 0

    for _, row in tqdm(qa_df.iterrows(), total=len(qa_df), desc="Processing"):
        task_id = row["task_id"]
        query = row["query"]
        ground_truth = str(row["ground_truth"])
        metric = row["metric"]
        task_desc = row["task_description"]

        # Get query embedding (already computed above)
        query_embedding = query_embedding_cache[query]
        task_desc_embedding = task_desc_embeddings[task_id]

        for llm_name in llm_names:
            cache_key = _query_hash(query, llm_name)

            if cache_key not in cache:
                missing_count += 1
                if skip_missing:
                    continue
                response = ""
                latency = 0.0
            else:
                cached = cache[cache_key]
                response = cached.get("response", "")
                latency = cached.get("latency", 0.0)

            # Evaluate response
            try:
                effect = engine.eval(response, ground_truth, metric)
            except Exception as e:
                logger.debug(f"Eval error for {task_id}/{llm_name}: {e}")
                effect = 0.0
                error_count += 1

            # Get cost (normalized)
            cost = llm_costs.get(llm_name, 0.5)

            records.append({
                "task_id": task_id,
                "task_description": task_desc,
                "task_description_embedding": str(task_desc_embedding),
                "query": query,
                "query_embedding": str(query_embedding),
                "ground_truth": ground_truth,
                "metric": metric,
                "llm": llm_name,
                "effect": effect,
                "cost": cost,
                "latency": latency,
            })

    logger.info(f"  Generated {len(records):,} records")
    logger.info(f"  Missing from cache: {missing_count:,}")
    logger.info(f"  Eval errors: {error_count:,}")

    # Create DataFrame and save
    router_df = pd.DataFrame.from_records(records)
    router_df.to_csv(output_path, index=False)
    logger.info(f"  Saved to {output_path}")

    # Summary stats
    logger.info("\n=== Summary ===")
    logger.info(f"Tasks: {router_df['task_id'].nunique()}")
    logger.info(f"Models: {router_df['llm'].nunique()}")
    logger.info(f"Total rows: {len(router_df):,}")

    logger.info("\nPer-model average effect:")
    model_effects = router_df.groupby("llm")["effect"].mean().sort_values(ascending=False)
    for model, effect in model_effects.items():
        logger.info(f"  {model:25s}: {effect:.3f}")

    return router_df


if __name__ == "__main__":
    regenerate_router_data()
