#!/usr/bin/env python3
"""Fix missing LLM responses in the response cache.

This script identifies queries missing LLM responses and re-queries those specific
(query, LLM) pairs. It's designed to fix data gaps from API failures without
regenerating the entire dataset.

Usage:
    python -m embergraphrouter.data_processing.fix_missing_llm_responses

    # Dry run (show what would be fixed without making API calls)
    python -m embergraphrouter.data_processing.fix_missing_llm_responses --dry-run

    # Fix specific tasks only
    python -m embergraphrouter.data_processing.fix_missing_llm_responses \
        --tasks mmlu_pro_law mmlu_pro_chemistry

    # Fix specific LLMs only
    python -m embergraphrouter.data_processing.fix_missing_llm_responses \
        --llms "GPT-5 Nano" "Gemini 3 Pro Preview"
"""

import argparse
import hashlib
import json
import logging
import signal
import sys
import threading
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import pandas as pd
from tqdm import tqdm

# Global flag for graceful shutdown
_shutdown_requested = False

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

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 cache key for query+llm pair."""
    return hashlib.sha256(f"{query}||{llm_name}".encode()).hexdigest()[:16]


class AdaptiveRateLimiter:
    """AIMD-style adaptive rate limiter per provider.

    Uses Additive Increase / Multiplicative Decrease (like TCP congestion control):
    - On success: gradually increase concurrency (+1 every `increase_after` successes)
    - On rate limit (429) or failure: cut concurrency in half
    - Bounded by min_concurrent and max_concurrent
    """

    def __init__(
        self,
        initial: int = 10,
        min_concurrent: int = 2,
        max_concurrent: int = 100,
        increase_after: int = 20,  # Increase by 1 after N consecutive successes
        decrease_factor: float = 0.5,  # Multiplicative decrease on failure
    ):
        self.current = initial
        self.min_concurrent = min_concurrent
        self.max_concurrent = max_concurrent
        self.increase_after = increase_after
        self.decrease_factor = decrease_factor

        self._semaphore = threading.Semaphore(initial)
        self._lock = threading.Lock()
        self._success_count = 0
        self._total_successes = 0
        self._total_failures = 0
        self._last_adjustment = time.time()

    def acquire(self):
        """Acquire a slot (blocks if at capacity)."""
        self._semaphore.acquire()

    def release_success(self):
        """Release slot after successful request - may increase capacity."""
        with self._lock:
            self._success_count += 1
            self._total_successes += 1

            # Additive increase: +1 capacity after N consecutive successes
            if self._success_count >= self.increase_after and self.current < self.max_concurrent:
                self.current += 1
                self._semaphore.release()  # Add one more slot
                self._semaphore.release()  # Release current
                self._success_count = 0
                logger.debug(f"Rate limit increased to {self.current}")
                return

        self._semaphore.release()

    def release_failure(self, is_rate_limit: bool = False):
        """Release slot after failed request - decreases capacity."""
        with self._lock:
            self._total_failures += 1
            self._success_count = 0  # Reset success streak

            # Multiplicative decrease on rate limit or repeated failures
            if is_rate_limit or self._total_failures % 3 == 0:
                new_limit = max(self.min_concurrent, int(self.current * self.decrease_factor))
                if new_limit < self.current:
                    # Drain excess capacity by not releasing
                    slots_to_drain = self.current - new_limit
                    self.current = new_limit
                    logger.warning(f"Rate limit decreased to {self.current} (draining {slots_to_drain} slots)")
                    # Don't release - let the semaphore drain naturally
                    return

        self._semaphore.release()

    def stats(self) -> dict:
        """Return current stats."""
        return {
            "current_limit": self.current,
            "total_successes": self._total_successes,
            "total_failures": self._total_failures,
            "success_rate": self._total_successes / max(1, self._total_successes + self._total_failures),
        }


# Adaptive rate limiters per provider
_rate_limiters: dict[str, AdaptiveRateLimiter] = {}


def _get_rate_limiter(provider: str) -> AdaptiveRateLimiter:
    """Get or create adaptive rate limiter for provider."""
    if provider not in _rate_limiters:
        # Initial limits based on typical provider quotas
        initial_limits = {
            "openai": 20,
            "google": 50,
            "anthropic": 30,
        }
        _rate_limiters[provider] = AdaptiveRateLimiter(
            initial=initial_limits.get(provider, 10),
            min_concurrent=2,
            max_concurrent=100,
        )
    return _rate_limiters[provider]


def _get_provider(llm_name: str) -> str:
    """Infer provider from LLM name."""
    name_lower = llm_name.lower()
    if "gpt" in name_lower or "codex" in name_lower:
        return "openai"
    elif "gemini" in name_lower:
        return "google"
    elif "claude" in name_lower:
        return "anthropic"
    return "openai"  # Default


def analyze_missing_responses(
    qa_path: Path,
    cache_path: Path,
    llm_path: Path,
    filter_tasks: list[str] | None = None,
    filter_llms: list[str] | None = None,
) -> tuple[list[tuple[str, str, str, str, str]], dict]:
    """Analyze which (query, llm) pairs are missing from cache.

    Args:
        qa_path: Path to unified_qa_data.csv
        cache_path: Path to response_cache.json
        llm_path: Path to LLM_Descriptions.json
        filter_tasks: If provided, only analyze these task_ids
        filter_llms: If provided, only analyze these LLM names

    Returns:
        Tuple of (missing_pairs, stats_dict)
        missing_pairs: List of (task_id, query, ground_truth, metric, llm_name)
    """
    # 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: {llm_names}")
    logger.info(f"  Cache entries: {len(cache):,}")

    # Apply filters
    if filter_tasks:
        qa_df = qa_df[qa_df["task_id"].isin(filter_tasks)]
        logger.info(f"  Filtered to tasks: {filter_tasks} ({len(qa_df):,} queries)")
    if filter_llms:
        llm_names = [n for n in llm_names if n in filter_llms]
        logger.info(f"  Filtered to LLMs: {llm_names}")

    # Find missing pairs
    logger.info("Scanning for missing cache entries...")
    missing_pairs = []
    stats = {
        "total_pairs": 0,
        "cached_pairs": 0,
        "cached_refusals": 0,  # Entries with refusal=True (content policy blocks)
        "missing_pairs": 0,
        "empty_responses": 0,
        "by_task": defaultdict(lambda: defaultdict(int)),
        "by_llm": defaultdict(int),
    }

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

        for llm_name in llm_names:
            stats["total_pairs"] += 1
            cache_key = _query_hash(query, llm_name)

            if cache_key not in cache:
                stats["missing_pairs"] += 1
                stats["by_task"][task_id][llm_name] += 1
                stats["by_llm"][llm_name] += 1
                missing_pairs.append((task_id, query, ground_truth, metric, llm_name))
            elif not cache[cache_key].get("response", "").strip():
                # Empty response - check if it's a known refusal
                if cache[cache_key].get("refusal"):
                    # Refusal is valid data - model declined to answer
                    stats["cached_pairs"] += 1
                    stats["cached_refusals"] += 1
                else:
                    # Unknown empty - needs fixing
                    stats["empty_responses"] += 1
                    stats["by_task"][task_id][llm_name] += 1
                    stats["by_llm"][llm_name] += 1
                    missing_pairs.append((task_id, query, ground_truth, metric, llm_name))
            else:
                stats["cached_pairs"] += 1

    return missing_pairs, stats


def fix_missing_responses(
    qa_path: Path,
    cache_path: Path,
    llm_path: Path,
    filter_tasks: list[str] | None = None,
    filter_llms: list[str] | None = None,
    max_workers: int = 20,
    dry_run: bool = False,
    save_interval: int = 50,
) -> None:
    """Fix missing LLM responses by re-querying APIs.

    Resumption: The script naturally resumes from where it left off because:
    1. Cache is saved periodically (every save_interval responses)
    2. On restart, it re-scans the cache and only processes still-missing pairs
    3. A checkpoint file tracks progress for visibility

    Args:
        qa_path: Path to unified_qa_data.csv
        cache_path: Path to response_cache.json
        llm_path: Path to LLM_Descriptions.json
        filter_tasks: If provided, only fix these task_ids
        filter_llms: If provided, only fix these LLM names
        max_workers: Max concurrent API calls
        dry_run: If True, analyze but don't make API calls
        save_interval: Save cache every N successful responses
    """
    # Checkpoint file for tracking progress across restarts
    checkpoint_path = cache_path.parent / "fix_checkpoint.json"

    # Load previous checkpoint if exists
    prev_fixed = 0
    if checkpoint_path.exists():
        try:
            with open(checkpoint_path) as f:
                ckpt = json.load(f)
                prev_fixed = ckpt.get("total_fixed", 0)
                logger.info(f"Resuming from checkpoint: {prev_fixed} previously fixed")
        except Exception:
            pass

    # Analyze missing pairs (this naturally excludes already-cached entries)
    missing_pairs, stats = analyze_missing_responses(
        qa_path, cache_path, llm_path, filter_tasks, filter_llms
    )

    # Report statistics
    logger.info("\n=== Missing Response Analysis ===")
    logger.info(f"Total pairs checked: {stats['total_pairs']:,}")
    logger.info(f"Cached (valid): {stats['cached_pairs']:,}")
    if stats["cached_refusals"] > 0:
        logger.info(f"  (includes {stats['cached_refusals']:,} content policy refusals)")
    logger.info(f"Missing from cache: {stats['missing_pairs']:,}")
    logger.info(f"Empty responses: {stats['empty_responses']:,}")
    logger.info(f"Total to fix: {len(missing_pairs):,}")

    if stats["by_llm"]:
        logger.info("\nMissing by LLM:")
        for llm, count in sorted(stats["by_llm"].items(), key=lambda x: -x[1]):
            logger.info(f"  {llm}: {count:,}")

    if stats["by_task"]:
        logger.info("\nMissing by task (top 10):")
        task_totals = {t: sum(d.values()) for t, d in stats["by_task"].items()}
        for task, count in sorted(task_totals.items(), key=lambda x: -x[1])[:10]:
            logger.info(f"  {task}: {count:,}")

    if dry_run:
        logger.info("\n[DRY RUN] No API calls made.")
        return

    if not missing_pairs:
        logger.info("\nNo missing responses to fix!")
        return

    # Load data for API calls
    llm_descriptions = loadjson(str(llm_path))
    llm_names = list(llm_descriptions.keys())
    engine = LLMEngine(llm_names, llm_descriptions)

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

    cache_lock = threading.Lock()
    save_lock = threading.Lock()
    completed = 0
    failures = 0
    refusals = 0  # Content policy refusals (empty responses with HTTP 200)
    session_start = time.time()

    def save_cache_and_checkpoint():
        """Save cache and checkpoint to disk for resumption."""
        with save_lock:
            # Save cache
            with open(cache_path, "w", encoding="utf-8") as f:
                json.dump(cache, f)
            # Save checkpoint
            with open(checkpoint_path, "w", encoding="utf-8") as f:
                json.dump({
                    "total_fixed": prev_fixed + completed,
                    "session_fixed": completed,
                    "session_refusals": refusals,
                    "session_failures": failures,
                    "session_elapsed_sec": time.time() - session_start,
                    "remaining": len(missing_pairs) - completed - failures - refusals,
                }, f, indent=2)

    def process_pair(pair: tuple, max_retries: int = 3) -> bool:
        """Process a single (query, llm) pair with adaptive rate limiting.

        Returns:
            True if we got a response (including refusals), False on hard failure.
        """
        nonlocal completed, failures, refusals
        task_id, query, ground_truth, metric, llm_name = pair

        llm_idx = llm_names.index(llm_name)
        provider = _get_provider(llm_name)
        limiter = _get_rate_limiter(provider)

        last_error = None
        is_rate_limit = False

        for attempt in range(max_retries):
            if _shutdown_requested:
                return False

            limiter.acquire()
            try:
                response, latency = engine.get_llm_response_with_latency(
                    query=query, llm_idx=llm_idx
                )

                if response and response.strip():
                    cache_key = _query_hash(query, llm_name)
                    with cache_lock:
                        cache[cache_key] = {"response": response, "latency": latency}
                        completed += 1

                        # Periodic save
                        if completed % save_interval == 0:
                            save_cache_and_checkpoint()
                            # Log adaptive limiter stats
                            stats = {p: _get_rate_limiter(p).stats() for p in ["openai", "google", "anthropic"]}
                            limits_str = ", ".join(f"{p}:{s['current_limit']}" for p, s in stats.items())
                            logger.info(f"  Checkpoint ({completed:,} fixed) | Limits: {limits_str}")

                    limiter.release_success()
                    return True
                else:
                    # Empty response with HTTP 200 = content policy refusal
                    # This is meaningful data - the model chose not to answer
                    # Store it and don't retry (retrying won't change the outcome)
                    cache_key = _query_hash(query, llm_name)
                    with cache_lock:
                        cache[cache_key] = {
                            "response": "",
                            "latency": latency,
                            "refusal": True,
                        }
                        refusals += 1
                    logger.info(f"⊘ {llm_name} refused (content policy)")
                    limiter.release_success()  # API worked, model just declined
                    return True  # Don't retry - refusal is a valid response

            except Exception as e:
                last_error = str(e)
                # Detect rate limit errors (429, "rate limit", "quota exceeded")
                error_lower = last_error.lower()
                is_rate_limit = "429" in last_error or "rate" in error_lower or "quota" in error_lower
                limiter.release_failure(is_rate_limit=is_rate_limit)

                if is_rate_limit:
                    # Extra backoff on rate limits
                    time.sleep(2 ** (attempt + 1))

            # Exponential backoff before retry (1s, 2s, 4s)
            if attempt < max_retries - 1:
                backoff = 2 ** attempt
                time.sleep(backoff)

        # All retries exhausted
        logger.warning(f"Failed {llm_name} after {max_retries} retries: {last_error}")
        failures += 1
        return False

    # Smart interleaving: sort pairs to maximize cross-provider parallelism
    # Group by provider, then interleave so consecutive requests hit different providers
    logger.info("Interleaving requests for optimal parallelism...")

    provider_buckets = defaultdict(list)
    for pair in missing_pairs:
        llm_name = pair[4]  # (task_id, query, ground_truth, metric, llm_name)
        provider = _get_provider(llm_name)
        provider_buckets[provider].append(pair)

    # Log distribution
    for provider, pairs in sorted(provider_buckets.items()):
        llm_counts = defaultdict(int)
        for p in pairs:
            llm_counts[p[4]] += 1
        logger.info(f"  {provider}: {len(pairs):,} requests ({dict(llm_counts)})")

    # Interleave: round-robin across providers
    interleaved_pairs = []
    provider_iters = {p: iter(pairs) for p, pairs in provider_buckets.items()}
    while provider_iters:
        exhausted = []
        for provider in list(provider_iters.keys()):
            try:
                interleaved_pairs.append(next(provider_iters[provider]))
            except StopIteration:
                exhausted.append(provider)
        for p in exhausted:
            del provider_iters[p]

    missing_pairs = interleaved_pairs

    # Process missing pairs
    logger.info(f"\n=== Fixing {len(missing_pairs):,} missing responses ===")
    logger.info(f"Using {max_workers} workers (per-provider rate limiting enabled)")
    logger.info("Press Ctrl+C to gracefully stop and save progress")

    # Set up signal handler for graceful shutdown
    def signal_handler(signum, frame):
        global _shutdown_requested
        if _shutdown_requested:
            logger.warning("Force quit requested. Saving and exiting...")
            save_cache_and_checkpoint()
            sys.exit(1)
        _shutdown_requested = True
        logger.warning("\nShutdown requested. Finishing current tasks and saving...")

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    start_time = time.time()

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(process_pair, p): p for p in missing_pairs}

        with tqdm(total=len(missing_pairs), desc="Fixing") as pbar:
            for future in as_completed(futures):
                if _shutdown_requested:
                    logger.info("Cancelling remaining tasks...")
                    for f in futures:
                        f.cancel()
                    break
                pbar.update(1)
                pbar.set_postfix({"fixed": completed, "refused": refusals, "failed": failures})

    # Final save
    save_cache_and_checkpoint()

    elapsed = time.time() - start_time
    logger.info("\n=== Fix Complete ===")
    logger.info(f"This session: {completed:,} fixed, {refusals:,} refusals, {failures:,} failed")
    logger.info(f"Total fixed (all sessions): {prev_fixed + completed:,}")
    if refusals > 0:
        logger.info("Refusals are content policy blocks - stored with refusal=True in cache")
    logger.info(f"Time: {elapsed/60:.1f} minutes")
    throughput = (completed + refusals) / max(1, elapsed)
    logger.info(f"Throughput: {throughput:.1f} requests/sec")
    logger.info(f"Cache saved to: {cache_path}")
    logger.info(f"Checkpoint saved to: {checkpoint_path}")

    # Log final adaptive rate limiter stats
    logger.info("\nAdaptive Rate Limiter Final Stats:")
    for provider in ["openai", "google", "anthropic"]:
        if provider in _rate_limiters:
            stats = _rate_limiters[provider].stats()
            logger.info(f"  {provider}: limit={stats['current_limit']}, "
                       f"success_rate={stats['success_rate']:.1%}, "
                       f"successes={stats['total_successes']}, failures={stats['total_failures']}")

    # Clean up checkpoint if all done
    remaining = len(missing_pairs) - completed - failures
    if remaining == 0 and failures == 0:
        checkpoint_path.unlink(missing_ok=True)
        logger.info("All missing responses fixed! Checkpoint removed.")


def main():
    parser = argparse.ArgumentParser(
        description="Fix missing LLM responses in cache"
    )
    parser.add_argument(
        "--qa-path",
        type=Path,
        default=_MODULE_DIR / "data" / "unified_qa_data.csv",
    )
    parser.add_argument(
        "--cache-path",
        type=Path,
        default=_MODULE_DIR / "data" / ".router_cache" / "response_cache.json",
    )
    parser.add_argument(
        "--llm-path",
        type=Path,
        default=_MODULE_DIR / "configs" / "LLM_Descriptions.json",
    )
    parser.add_argument(
        "--tasks",
        nargs="+",
        help="Only fix these task IDs",
    )
    parser.add_argument(
        "--llms",
        nargs="+",
        help="Only fix these LLM names",
    )
    parser.add_argument(
        "--max-workers",
        type=int,
        default=20,
        help="Max concurrent API calls",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Analyze only, don't make API calls",
    )
    parser.add_argument(
        "--save-interval",
        type=int,
        default=10,
        help="Save cache every N responses (default: 10 for frequent checkpointing)",
    )

    args = parser.parse_args()

    fix_missing_responses(
        qa_path=args.qa_path,
        cache_path=args.cache_path,
        llm_path=args.llm_path,
        filter_tasks=args.tasks,
        filter_llms=args.llms,
        max_workers=args.max_workers,
        dry_run=args.dry_run,
        save_interval=args.save_interval,
    )


if __name__ == "__main__":
    main()
