#!/usr/bin/env python3
"""Re-evaluate router data using cached responses with fixed extraction logic.

This script reads cached LLM responses and re-runs evaluation without making
any API calls. Use this after fixing extraction bugs or metrics.

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

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

# Import evaluation logic
try:
    from .llm_engine import LLMEngine
    from .utils import get_embedding, loadjson, savepkl
except ImportError:
    from llm_engine import LLMEngine
    from utils import get_embedding, loadjson, savepkl


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 reevaluate_router_data(
    qa_path: Path,
    llm_path: Path,
    cache_path: Path,
    output_path: Path,
    llm_embedding_path: Path,
    existing_router_data_path: Path | None = None,
) -> None:
    """Re-evaluate all cached responses with current extraction/evaluation logic.

    This function ONLY recalculates effect scores using the fixed evaluation logic.
    Costs, latencies, and embeddings are preserved from the existing router data
    to avoid expensive API calls for token counting.

    Args:
        qa_path: Path to unified_qa_data.csv
        llm_path: Path to LLM_Descriptions.json
        cache_path: Path to response_cache.json
        output_path: Path to write updated router_data.csv
        llm_embedding_path: Path to write LLM embeddings
        existing_router_data_path: Path to existing router_data.csv (for costs/embeddings)
    """
    # Load data
    logger.info("Loading data sources...")
    qa_data = pd.read_csv(qa_path)
    llm_description = loadjson(str(llm_path))
    llm_names = list(llm_description.keys())

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

    # Load existing router data for costs, latencies, embeddings.
    # Auto-detect format: Parquet files start with b'PAR1' magic bytes.
    existing_df = None
    if existing_router_data_path and existing_router_data_path.exists():
        with open(existing_router_data_path, "rb") as fp:
            magic = fp.read(4)
        if magic == b"PAR1":
            existing_df = pd.read_parquet(existing_router_data_path)
        else:
            existing_df = pd.read_csv(existing_router_data_path)
        logger.info(f"  Loaded existing router data: {len(existing_df)} rows")
        # Create lookup key with task_id to handle same query in different tasks
        existing_df['_lookup'] = (
            existing_df['task_id'] + '||' + existing_df['query'] + '||' + existing_df['llm']
        )
        # Handle duplicates by keeping first occurrence
        existing_df = existing_df.drop_duplicates(subset='_lookup', keep='first')
        existing_lookup = existing_df.set_index('_lookup').to_dict('index')
    else:
        existing_lookup = {}
        logger.warning("  No existing router data - costs will be set to 0")

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

    # Tasks to exclude from router training data.
    #   alpaca_data — instruction-following, no correct answer exists
    #   multi_news  — abstractive summarization, ROUGE is continuous
    EXCLUDED_TASKS = {"alpaca_data", "multi_news"}

    # Metric overrides: remap from inappropriate to official evaluation.
    #   cruxeval_output: em → cruxeval_em (ast.literal_eval comparison)
    #   codexglue_defect: em → binary_em (extracts 0/1 from verbose responses)
    METRIC_OVERRIDES = {
        "cruxeval_output": "cruxeval_em",
        "codexglue_defect": "binary_em",
    }

    before_len = len(qa_data)
    qa_data = qa_data[~qa_data['task_id'].isin(EXCLUDED_TASKS)]
    dropped = before_len - len(qa_data)
    if dropped > 0:
        logger.info(
            f"  Dropped {dropped} queries from excluded tasks: "
            f"{EXCLUDED_TASKS}"
        )

    # Initialize evaluation engine (only for eval, not cost)
    engine = LLMEngine(llm_names=llm_names, llm_description=llm_description)

    # Process each query
    records = []
    missing_cache = 0
    total_pairs = len(qa_data) * len(llm_names)
    effects_changed = 0

    logger.info(f"Re-evaluating {total_pairs} query-LLM pairs (effect only, no API calls)...")

    for idx, row in qa_data.iterrows():
        query = row['query']
        task_id = row['task_id']
        task_description = row['task_description']
        ground_truth = row['ground_truth']
        metric = METRIC_OVERRIDES.get(task_id, row['metric'])

        for llm_idx, llm_name in enumerate(llm_names):
            cache_key = _query_hash(query, llm_name)
            lookup_key = task_id + '||' + query + '||' + llm_name

            if cache_key not in cache:
                missing_cache += 1
                continue

            cached = cache[cache_key]
            response = cached.get('response', '')
            latency = cached.get('latency', 0.0)

            # Re-evaluate with fixed logic
            new_effect = engine.eval(
                prediction=response,
                ground_truth=ground_truth,
                metric=metric,
            )

            # Get existing data (cost, embeddings) or defaults
            existing = existing_lookup.get(lookup_key, {})
            old_effect = existing.get('effect', 0.0)
            cost = existing.get('cost', 0.0)  # Preserve original cost

            # Use existing embeddings if available
            query_embedding = existing.get('query_embedding', '')
            task_embedding = existing.get('task_description_embedding', '')

            if new_effect != old_effect:
                effects_changed += 1

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

        # Progress logging
        if (idx + 1) % 2000 == 0:
            logger.info(f"  Processed {idx + 1}/{len(qa_data)} queries...")

    if missing_cache > 0:
        logger.warning(f"Missing cache entries: {missing_cache}")

    # Create DataFrame and normalize
    logger.info("Normalizing scores...")
    df = pd.DataFrame.from_records(records)

    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)
    if "latency" in df.columns:
        df["latency"] = df.groupby(["task_id", "llm"])["latency"].transform(
            lambda s: s.clip(upper=s.quantile(0.99))
        )
        df["latency"] = df.groupby("task_id")["latency"].transform(_minmax)

    # Save outputs
    output_path.parent.mkdir(parents=True, exist_ok=True)
    if str(output_path).endswith(".parquet"):
        df.to_parquet(output_path, index=False)
    else:
        df.to_csv(output_path, index=False)
    logger.info(f"Saved: {output_path}")

    # Save LLM embeddings only if they don't already exist.
    # Qwen3 embeddings are generated separately by generate_v2_embeddings.py
    # and should not be overwritten with the default provider (Gemini).
    if not llm_embedding_path.exists():
        all_llm_descriptions = [
            llm_description[name]["feature"] for name in llm_names
        ]
        llm_embeddings = get_embedding(all_llm_descriptions)
        savepkl(llm_embeddings, str(llm_embedding_path))
        logger.info(f"Generated and saved: {llm_embedding_path}")
    else:
        logger.info(f"Kept existing LLM embeddings: {llm_embedding_path}")

    # Save metadata
    task_edge_counts = df.groupby("task_id").size()
    task_query_counts = {
        task_id: int(count // len(llm_names))
        for task_id, count in task_edge_counts.items()
    }

    metadata = {
        "llm_names": llm_names,
        "total_rows": int(len(df)),
        "total_queries": int(len(df) // len(llm_names)),
        "task_query_counts": task_query_counts,
    }

    metadata_path = output_path.with_suffix(output_path.suffix + ".metadata.json")
    metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
    logger.info(f"Saved: {metadata_path}")

    # Summary statistics
    logger.info("=" * 60)
    logger.info("RE-EVALUATION COMPLETE")
    logger.info("=" * 60)
    logger.info(f"Total records: {len(df)}")
    logger.info(f"Total queries: {len(df) // len(llm_names)}")
    logger.info(f"Effects changed: {effects_changed} ({effects_changed/len(df)*100:.1f}%)")

    logger.info("\nPer-LLM effect averages:")
    for llm in llm_names:
        llm_df = df[df['llm'] == llm]
        avg = llm_df['effect'].mean()
        zero_pct = (llm_df['effect'] == 0).mean() * 100
        logger.info(f"  {llm:25s}: avg={avg:.3f}, zero={zero_pct:.1f}%")

    logger.info("\nPer-metric effect averages:")
    for metric in df['metric'].unique():
        m_df = df[df['metric'] == metric]
        avg = m_df['effect'].mean()
        zero_pct = (m_df['effect'] == 0).mean() * 100
        logger.info(f"  {metric:15s}: avg={avg:.3f}, zero={zero_pct:.1f}%")


if __name__ == "__main__":
    import argparse
    import os
    import sys

    sys.path.insert(0, str(Path(__file__).resolve().parent))
    os.environ["KMP_DUPLICATE_LIB_OK"] = "True"

    parser = argparse.ArgumentParser(description="Re-evaluate from cache")
    parser.add_argument(
        "--config",
        default="config_v2_2026.yaml",
        help="Config filename in configs/ dir (default: config_v2_2026.yaml)",
    )
    cli_args = parser.parse_args()

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

    # Resolve paths
    for key in (
        "unified_qa_data_path",
        "llm_description_path",
        "saved_router_data_path",
        "llm_embedding_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"
    output_path = Path(config["saved_router_data_path"])

    reevaluate_router_data(
        qa_path=Path(config["unified_qa_data_path"]),
        llm_path=Path(config["llm_description_path"]),
        cache_path=cache_dir / "response_cache.json",
        output_path=output_path,
        llm_embedding_path=Path(config["llm_embedding_path"]),
        existing_router_data_path=output_path,  # Reuse existing costs/embeddings
    )
