#!/usr/bin/env python3
"""Fix BBEH task metrics and re-evaluate cached responses.

This script:
1. Updates unified_qa_data.csv to use 'bbeh' metric for all BBEH tasks
2. Re-evaluates cached responses with the corrected metric
3. Regenerates router_data.csv with proper effect scores

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

import hashlib
import json
import logging
from pathlib import Path

import pandas as pd

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 loadjson
except ImportError:
    from llm_engine import LLMEngine
    from utils import loadjson


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 fix_bbeh_metrics_and_reevaluate(
    qa_path: Path,
    router_data_path: Path,
    cache_path: Path,
    llm_path: Path,
    output_qa_path: Path | None = None,
    output_router_path: Path | None = None,
) -> None:
    """Fix BBEH metrics and re-evaluate all cached responses.

    Args:
        qa_path: Path to unified_qa_data.csv
        router_data_path: Path to existing router_data.csv
        cache_path: Path to response_cache.json
        llm_path: Path to LLM_Descriptions.json
        output_qa_path: Path to write updated unified_qa_data.csv (default: overwrite)
        output_router_path: Path to write updated router_data.csv (default: overwrite)
    """
    if output_qa_path is None:
        output_qa_path = qa_path
    if output_router_path is None:
        output_router_path = router_data_path

    # Load data
    logger.info("Loading data sources...")
    qa_df = pd.read_csv(qa_path)
    router_df = pd.read_csv(router_data_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)

    logger.info(f"  QA data: {len(qa_df)} queries")
    logger.info(f"  Router data: {len(router_df)} rows")
    logger.info(f"  Cache entries: {len(cache)}")

    # Step 1: Update metrics in unified_qa_data.csv
    logger.info("\n=== Step 1: Updating BBEH metrics in unified_qa_data.csv ===")
    bbeh_tasks = qa_df[qa_df["task_id"].str.startswith("bbeh_")]["task_id"].unique()
    logger.info(f"BBEH tasks found: {list(bbeh_tasks)}")

    old_metrics = {}
    for task in bbeh_tasks:
        old_metric = qa_df[qa_df["task_id"] == task]["metric"].iloc[0]
        old_metrics[task] = old_metric
        logger.info(f"  {task}: {old_metric} -> bbeh")

    # Update metrics for BBEH tasks
    qa_df.loc[qa_df["task_id"].str.startswith("bbeh_"), "metric"] = "bbeh"

    # Save updated QA data
    qa_df.to_csv(output_qa_path, index=False)
    logger.info(f"  Saved updated QA data to {output_qa_path}")

    # Step 2: Re-evaluate BBEH tasks in router_data.csv
    logger.info("\n=== Step 2: Re-evaluating BBEH tasks in router_data.csv ===")

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

    # Track changes
    effects_changed = 0
    effects_improved = 0
    effects_worsened = 0
    missing_cache = 0

    # Process only BBEH rows
    bbeh_mask = router_df["task_id"].str.startswith("bbeh_")
    bbeh_indices = router_df[bbeh_mask].index
    total_bbeh = len(bbeh_indices)

    logger.info(f"Re-evaluating {total_bbeh} BBEH rows...")

    for i, idx in enumerate(bbeh_indices):
        row = router_df.loc[idx]
        query = row["query"]
        llm_name = row["llm"]
        ground_truth = row["ground_truth"]
        old_effect = row["effect"]

        # Get cached response
        cache_key = _query_hash(query, llm_name)
        if cache_key not in cache:
            missing_cache += 1
            continue

        response = cache[cache_key].get("response", "")

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

        # Update if changed
        if new_effect != old_effect:
            effects_changed += 1
            if new_effect > old_effect:
                effects_improved += 1
            else:
                effects_worsened += 1
            router_df.at[idx, "effect"] = new_effect

        # Progress logging
        if (i + 1) % 1000 == 0:
            logger.info(f"  Progress: {i+1}/{total_bbeh}")

    logger.info(f"\n=== Re-evaluation Results ===")
    logger.info(f"  Total BBEH rows: {total_bbeh}")
    logger.info(f"  Missing cache: {missing_cache}")
    logger.info(f"  Effects changed: {effects_changed}")
    logger.info(f"  Effects improved (0->1): {effects_improved}")
    logger.info(f"  Effects worsened (1->0): {effects_worsened}")

    # Step 3: Update metric column in router_data if it exists
    if "metric" in router_df.columns:
        router_df.loc[bbeh_mask, "metric"] = "bbeh"

    # Save updated router data
    router_df.to_csv(output_router_path, index=False)
    logger.info(f"\n  Saved updated router data to {output_router_path}")

    # Step 4: Print per-model comparison
    logger.info("\n=== Per-Model Effect Comparison (BBEH tasks only) ===")
    logger.info("Model                    | Old Avg | New Avg | Change")
    logger.info("-" * 55)

    # Reload to get fresh stats
    new_router_df = pd.read_csv(output_router_path)
    old_router_df = pd.read_csv(router_data_path) if output_router_path != router_data_path else router_df

    for llm in llm_names:
        old_bbeh = old_router_df[
            (old_router_df["llm"] == llm) &
            (old_router_df["task_id"].str.startswith("bbeh_"))
        ]["effect"].mean()
        new_bbeh = new_router_df[
            (new_router_df["llm"] == llm) &
            (new_router_df["task_id"].str.startswith("bbeh_"))
        ]["effect"].mean()
        change = new_bbeh - old_bbeh
        sign = "+" if change > 0 else ""
        logger.info(f"{llm:24} | {old_bbeh:.4f} | {new_bbeh:.4f} | {sign}{change:.4f}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Fix BBEH metrics and re-evaluate")
    parser.add_argument(
        "--qa-path",
        type=Path,
        default=_MODULE_DIR / "data" / "unified_qa_data.csv",
        help="Path to unified_qa_data.csv",
    )
    parser.add_argument(
        "--router-path",
        type=Path,
        default=_MODULE_DIR / "data" / "router_data.csv",
        help="Path to router_data.csv",
    )
    parser.add_argument(
        "--cache-path",
        type=Path,
        default=_MODULE_DIR / "data" / ".router_cache" / "response_cache.json",
        help="Path to response_cache.json",
    )
    parser.add_argument(
        "--llm-path",
        type=Path,
        default=_MODULE_DIR / "configs" / "LLM_Descriptions.json",
        help="Path to LLM_Descriptions.json",
    )
    parser.add_argument(
        "--output-qa",
        type=Path,
        default=None,
        help="Output path for updated QA data (default: overwrite input)",
    )
    parser.add_argument(
        "--output-router",
        type=Path,
        default=None,
        help="Output path for updated router data (default: overwrite input)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Don't write output files, just show what would change",
    )

    args = parser.parse_args()

    if args.dry_run:
        # For dry run, write to temp files
        import tempfile
        args.output_qa = Path(tempfile.mktemp(suffix="_qa.csv"))
        args.output_router = Path(tempfile.mktemp(suffix="_router.csv"))
        logger.info("DRY RUN: Writing to temp files")

    fix_bbeh_metrics_and_reevaluate(
        qa_path=args.qa_path,
        router_data_path=args.router_path,
        cache_path=args.cache_path,
        llm_path=args.llm_path,
        output_qa_path=args.output_qa,
        output_router_path=args.output_router,
    )
