"""Cost calibration script for router data.

Samples queries from each task type, runs them through all LLMs via Ember,
and captures actual token usage for building a cost prediction model.

Features:
    - Smart parallelism: concurrent calls per provider to respect rate limits
    - Resumption: skips already-completed (task_id, llm) pairs on restart
    - Incremental saves: writes results after each batch

Usage:
    python scripts/cost_calibration.py --samples-per-task 5 --output data/cost_calibration.csv
    python scripts/cost_calibration.py --resume --output data/cost_calibration.csv  # Resume
"""

import argparse
import hashlib
import logging
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any

import pandas as pd

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

# LLM name to Ember model ID mapping
LLM_TO_MODEL = {
    "Gemini 2.5 Flash Lite": "gemini-2.5-flash-lite",
    "Claude Haiku 4.5": "claude-haiku-4-5-20251001",
    "GPT-5 Nano": "gpt-5-nano",
    "GPT-5.1": "gpt-5.1",
    "GPT-5.1 Codex": "gpt-5.1-codex",
    "Claude Opus 4.5": "claude-opus-4-5-20251101",
    "Gemini 3 Flash Preview": "gemini-3-flash-preview",
    "Gemini 3 Pro Preview": "gemini-3-pro-preview",
}

# Provider groupings for rate-limit-aware parallelism
LLM_PROVIDERS = {
    "openai": ["GPT-5 Nano", "GPT-5.1", "GPT-5.1 Codex"],
    "anthropic": ["Claude Haiku 4.5", "Claude Opus 4.5"],
    "google": ["Gemini 2.5 Flash Lite", "Gemini 3 Flash Preview", "Gemini 3 Pro Preview"],
}

# Pricing from LLM_Descriptions.json ($ per million tokens)
LLM_PRICING = {
    "Gemini 2.5 Flash Lite": {"input": 0.1, "output": 0.4},
    "Claude Haiku 4.5": {"input": 1.0, "output": 5.0},
    "GPT-5 Nano": {"input": 0.05, "output": 0.4},
    "GPT-5.1": {"input": 1.25, "output": 10.0},
    "GPT-5.1 Codex": {"input": 1.25, "output": 10.0},
    "Claude Opus 4.5": {"input": 5.0, "output": 25.0},
    "Gemini 3 Flash Preview": {"input": 0.5, "output": 3.0},
    "Gemini 3 Pro Preview": {"input": 2.0, "output": 12.0},
}


def query_hash(task_id: str, query: str) -> str:
    """Create a short hash to identify a query for resumption."""
    return hashlib.md5(f"{task_id}:{query[:100]}".encode()).hexdigest()[:12]


def sample_queries(df: pd.DataFrame, samples_per_task: int, seed: int = 42) -> pd.DataFrame:
    """Sample diverse queries from each task type."""
    sampled = []
    for task_id in df["task_id"].unique():
        task_queries = df[df["task_id"] == task_id]["query"].drop_duplicates()
        n_samples = min(samples_per_task, len(task_queries))
        task_sample = task_queries.sample(n=n_samples, random_state=seed)
        for query in task_sample:
            sampled.append({
                "task_id": task_id,
                "query": query,
                "query_hash": query_hash(task_id, query),
            })
    return pd.DataFrame(sampled)


def load_completed(output_path: Path) -> set[tuple[str, str]]:
    """Load already-completed (query_hash, llm) pairs from output file.

    Only counts successful calls as completed - failed calls will be retried.
    """
    if not output_path.exists():
        return set()
    try:
        df = pd.read_csv(output_path)
        # Only count successful calls as completed
        success_df = df[df["success"] == True]
        if "query_hash" not in df.columns:
            # Legacy format without query_hash - rebuild from task_id + query
            completed = set()
            for _, row in success_df.iterrows():
                qh = query_hash(row["task_id"], row.get("query", "")[:100])
                completed.add((qh, row["llm"]))
            return completed
        return {(row["query_hash"], row["llm"]) for _, row in success_df.iterrows()}
    except Exception as e:
        logger.warning(f"Could not load existing results: {e}")
        return set()


def call_single_llm(
    task_id: str,
    query: str,
    qhash: str,
    llm_name: str,
    max_output_tokens: int = 1024,
) -> dict[str, Any]:
    """Make a single API call to one LLM. Returns result dict."""
    import ember

    model_id = LLM_TO_MODEL.get(llm_name)
    if not model_id:
        return {
            "task_id": task_id,
            "query": query[:200],
            "query_hash": qhash,
            "query_len": len(query),
            "llm": llm_name,
            "model_id": "",
            "success": False,
            "error": f"No model mapping for {llm_name}",
        }

    try:
        start_time = time.time()
        response = ember.models.responses(
            model=model_id,
            input=query,
            max_output_tokens=max_output_tokens,
        )
        latency = time.time() - start_time

        usage = response.usage or {}
        result = {
            "task_id": task_id,
            "query": query[:200],
            "query_hash": qhash,
            "query_len": len(query),
            "llm": llm_name,
            "model_id": model_id,
            "input_tokens": usage.get("prompt_tokens", 0),
            "output_tokens": usage.get("completion_tokens", 0),
            "total_tokens": usage.get("total_tokens", 0),
            "cost_usd": usage.get("cost", 0),
            "latency_seconds": latency,
            "response_len": len(response.text) if response.text else 0,
            "success": True,
            "error": None,
        }

        # Extract details if available
        details = usage.get("details", {})
        if details:
            result["reasoning_tokens"] = details.get("reasoning_tokens", 0)
            result["cache_read_tokens"] = details.get("cache_read_input_tokens", 0)
            result["cache_creation_tokens"] = details.get("cache_creation_input_tokens", 0)

        return result

    except Exception as e:
        return {
            "task_id": task_id,
            "query": query[:200],
            "query_hash": qhash,
            "query_len": len(query),
            "llm": llm_name,
            "model_id": model_id,
            "input_tokens": 0,
            "output_tokens": 0,
            "total_tokens": 0,
            "cost_usd": 0,
            "latency_seconds": 0,
            "response_len": 0,
            "success": False,
            "error": str(e)[:200],
        }


def run_calibration_parallel(
    samples: pd.DataFrame,
    llms: list[str],
    completed: set[tuple[str, str]],
    output_path: Path,
    workers_per_provider: int = 2,
    delay_between_calls: float = 0.3,
    max_output_tokens: int = 1024,
) -> list[dict]:
    """Run calibration with smart parallelism by provider.

    Args:
        samples: DataFrame with task_id, query, query_hash columns
        llms: List of LLM names to calibrate
        completed: Set of (query_hash, llm) pairs already done
        output_path: Path to save incremental results
        workers_per_provider: Max concurrent calls per provider
        delay_between_calls: Delay between submitting tasks
        max_output_tokens: Max tokens for responses

    Returns:
        List of result dicts
    """
    # Build work items, skipping completed ones
    work_items = []
    for _, row in samples.iterrows():
        for llm_name in llms:
            if (row["query_hash"], llm_name) in completed:
                continue
            work_items.append({
                "task_id": row["task_id"],
                "query": row["query"],
                "query_hash": row["query_hash"],
                "llm_name": llm_name,
            })

    if not work_items:
        logger.info("All work items already completed!")
        return []

    total_work = len(work_items)
    skipped = len(samples) * len(llms) - total_work
    logger.info(f"Total work: {total_work} calls ({skipped} skipped from previous run)")

    results = []
    completed_count = 0

    # Process in batches by provider to respect rate limits
    # Use ThreadPoolExecutor with limited workers per provider
    with ThreadPoolExecutor(max_workers=workers_per_provider * len(LLM_PROVIDERS)) as executor:
        futures = {}

        for item in work_items:
            future = executor.submit(
                call_single_llm,
                item["task_id"],
                item["query"],
                item["query_hash"],
                item["llm_name"],
                max_output_tokens,
            )
            futures[future] = item
            time.sleep(delay_between_calls)  # Stagger submissions

        # Collect results as they complete
        for future in as_completed(futures):
            item = futures[future]
            try:
                result = future.result(timeout=120)
                results.append(result)
                completed_count += 1

                status = "OK" if result["success"] else f"FAIL: {result.get('error', '')[:50]}"
                logger.info(
                    f"[{completed_count}/{total_work}] {item['llm_name']} on "
                    f"{item['task_id'][:20]}... {status}"
                )

                # Incremental save every 10 results
                if completed_count % 10 == 0:
                    _save_incremental(results, output_path, completed)

            except Exception as e:
                logger.error(f"Future failed for {item['llm_name']}: {e}")
                results.append({
                    "task_id": item["task_id"],
                    "query": item["query"][:200],
                    "query_hash": item["query_hash"],
                    "query_len": len(item["query"]),
                    "llm": item["llm_name"],
                    "model_id": LLM_TO_MODEL.get(item["llm_name"], ""),
                    "success": False,
                    "error": f"Future timeout: {e}"[:200],
                })

    return results


def _save_incremental(
    new_results: list[dict],
    output_path: Path,
    existing_completed: set[tuple[str, str]],
) -> None:
    """Save new results, merging with existing file."""
    new_df = pd.DataFrame(new_results)

    if output_path.exists() and existing_completed:
        try:
            existing_df = pd.read_csv(output_path)
            combined = pd.concat([existing_df, new_df], ignore_index=True)
            combined.to_csv(output_path, index=False)
        except Exception:
            new_df.to_csv(output_path, index=False)
    else:
        new_df.to_csv(output_path, index=False)


def run_calibration_sequential(
    samples: pd.DataFrame,
    llms: list[str],
    completed: set[tuple[str, str]],
    delay_between_calls: float = 0.5,
    max_output_tokens: int = 1024,
) -> list[dict]:
    """Run calibration sequentially (original behavior)."""
    results = []
    total_calls = len(samples) * len(llms)
    call_count = 0
    skipped = 0

    for _, row in samples.iterrows():
        task_id = row["task_id"]
        query = row["query"]
        qhash = row["query_hash"]

        for llm_name in llms:
            call_count += 1

            # Skip if already completed
            if (qhash, llm_name) in completed:
                skipped += 1
                continue

            logger.info(f"[{call_count}/{total_calls}] {llm_name} on {task_id[:20]}...")
            result = call_single_llm(task_id, query, qhash, llm_name, max_output_tokens)
            results.append(result)
            time.sleep(delay_between_calls)

    if skipped > 0:
        logger.info(f"Skipped {skipped} already-completed calls")

    return results


def build_cost_model(calibration_df: pd.DataFrame) -> dict:
    """Build a simple cost prediction model from calibration data.

    Returns median output tokens per (task_id, llm) pair.
    """
    # Filter to successful calls
    success_df = calibration_df[calibration_df["success"] == True]

    # Calculate medians per (task, llm)
    model = {}
    for (task_id, llm), group in success_df.groupby(["task_id", "llm"]):
        model[(task_id, llm)] = {
            "median_input_tokens": group["input_tokens"].median(),
            "median_output_tokens": group["output_tokens"].median(),
            "median_cost": group["cost_usd"].median(),
            "input_to_output_ratio": (
                group["output_tokens"].median() / group["input_tokens"].median()
                if group["input_tokens"].median() > 0 else 1.0
            ),
            "samples": len(group),
        }

    # Also calculate global per-task and per-llm averages as fallbacks
    for task_id, group in success_df.groupby("task_id"):
        model[("__task__", task_id)] = {
            "median_output_tokens": group["output_tokens"].median(),
            "input_to_output_ratio": (
                group["output_tokens"].median() / group["input_tokens"].median()
                if group["input_tokens"].median() > 0 else 1.0
            ),
        }

    for llm, group in success_df.groupby("llm"):
        model[("__llm__", llm)] = {
            "median_output_tokens": group["output_tokens"].median(),
            "input_to_output_ratio": (
                group["output_tokens"].median() / group["input_tokens"].median()
                if group["input_tokens"].median() > 0 else 1.0
            ),
        }

    return model


def estimate_cost(
    query_len: int,
    task_id: str,
    llm: str,
    cost_model: dict,
    chars_per_token: float = 4.0,
) -> tuple[int, int, float]:
    """Estimate input tokens, output tokens, and cost for a query.

    Returns:
        Tuple of (input_tokens, output_tokens, cost_usd)
    """
    # Estimate input tokens from query length
    input_tokens = int(query_len / chars_per_token)

    # Get output prediction from model
    key = (task_id, llm)
    if key in cost_model:
        ratio = cost_model[key]["input_to_output_ratio"]
    elif ("__task__", task_id) in cost_model:
        ratio = cost_model[("__task__", task_id)]["input_to_output_ratio"]
    elif ("__llm__", llm) in cost_model:
        ratio = cost_model[("__llm__", llm)]["input_to_output_ratio"]
    else:
        ratio = 2.0  # Default fallback: output ~2x input

    output_tokens = int(input_tokens * ratio)

    # Calculate cost from pricing
    pricing = LLM_PRICING.get(llm, {"input": 1.0, "output": 5.0})
    cost_usd = (
        input_tokens * pricing["input"] / 1_000_000
        + output_tokens * pricing["output"] / 1_000_000
    )

    return input_tokens, output_tokens, cost_usd


def main():
    parser = argparse.ArgumentParser(description="Cost calibration for router data")
    parser.add_argument(
        "--input",
        type=str,
        default="data/router_data.csv",
        help="Path to router data CSV",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="data/cost_calibration.csv",
        help="Path to save calibration results",
    )
    parser.add_argument(
        "--samples-per-task",
        type=int,
        default=5,
        help="Number of queries to sample per task",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=0.3,
        help="Delay between API calls in seconds",
    )
    parser.add_argument(
        "--parallel",
        action="store_true",
        help="Use parallel execution (default: sequential)",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=2,
        help="Workers per provider for parallel mode",
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Resume from existing output file, skipping completed calls",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Only sample queries, don't run API calls",
    )
    args = parser.parse_args()

    output_path = Path(args.output)

    # Load data
    logger.info(f"Loading data from {args.input}")
    df = pd.read_csv(args.input)
    logger.info(f"Loaded {len(df):,} rows, {df['task_id'].nunique()} tasks")

    # Sample queries
    samples = sample_queries(df, args.samples_per_task)
    logger.info(f"Sampled {len(samples)} queries across {samples['task_id'].nunique()} tasks")

    if args.dry_run:
        logger.info("Dry run - showing sample queries:")
        print(samples.head(20).to_string())
        return

    # Load completed calls for resumption
    completed = set()
    if args.resume or output_path.exists():
        completed = load_completed(output_path)
        if completed:
            logger.info(f"Loaded {len(completed)} completed calls for resumption")

    # Run calibration
    llms = list(LLM_TO_MODEL.keys())
    logger.info(f"Running calibration with {len(llms)} LLMs...")
    total_calls = len(samples) * len(llms)
    remaining = total_calls - len(completed)
    logger.info(f"Total API calls: {total_calls} ({remaining} remaining)")

    if args.parallel:
        logger.info(f"Parallel mode: {args.workers} workers per provider")
        results = run_calibration_parallel(
            samples,
            llms,
            completed,
            output_path,
            workers_per_provider=args.workers,
            delay_between_calls=args.delay,
        )
    else:
        results = run_calibration_sequential(
            samples,
            llms,
            completed,
            delay_between_calls=args.delay,
        )

    if not results:
        logger.info("No new results to save")
        if output_path.exists():
            results_df = pd.read_csv(output_path)
        else:
            return
    else:
        # Save final results
        new_df = pd.DataFrame(results)
        if output_path.exists():
            existing_df = pd.read_csv(output_path)
            # Remove old failed rows that we just retried
            new_keys = {(row["query_hash"], row["llm"]) for _, row in new_df.iterrows()}
            existing_df = existing_df[
                ~existing_df.apply(
                    lambda r: (r.get("query_hash", query_hash(r["task_id"], r.get("query", "")[:100])), r["llm"]) in new_keys,
                    axis=1
                )
            ]
            results_df = pd.concat([existing_df, new_df], ignore_index=True)
        else:
            results_df = new_df

        results_df.to_csv(output_path, index=False)
        logger.info(f"Saved {len(results)} new results to {output_path}")

    # Build and display cost model
    cost_model = build_cost_model(results_df)

    logger.info("\n=== Cost Model Summary ===")
    success_rate = results_df["success"].mean() * 100
    logger.info(f"Success rate: {success_rate:.1f}%")
    logger.info(f"Total records: {len(results_df)}")

    logger.info("\nMedian output tokens per LLM:")
    for llm in llms:
        llm_data = results_df[(results_df["llm"] == llm) & (results_df["success"])]
        if len(llm_data) > 0:
            logger.info(
                f"  {llm}: input={llm_data['input_tokens'].median():.0f}, "
                f"output={llm_data['output_tokens'].median():.0f}, "
                f"cost=${llm_data['cost_usd'].median():.6f}"
            )


if __name__ == "__main__":
    main()
