#!/usr/bin/env python3
"""Multi-sample router data collection with variance-aware aggregation.

This script extends construct_router_data.py with multi-sample collection based on
variance analysis findings:
- 23.6% of (query, LLM) pairs have high variance (σ > 0.1)
- Gemini Flash Lite: 45% high variance
- GPT-5.1: 33% high variance
- Claude Opus 4.5: 12% high variance (most consistent)
- GPT-5.1 Codex: 0% variance (deterministic, no temperature support)

Key features:
- Selective multi-sampling: more samples for high-variance LLMs
- Temperature-controlled sampling (0.7 default)
- Aggregated statistics: mean, std, min, max, n_samples
- AIMD-style adaptive rate limiting
- Checkpoint-based resumption
- Thread-safe parallel processing

Usage:
    python -m embergraphrouter.data_processing.construct_router_data_multisample

    # Custom parameters
    python -m embergraphrouter.data_processing.construct_router_data_multisample \
        --n-samples 5 --temperature 0.7 --max-workers 20

    # Resume interrupted run
    python -m embergraphrouter.data_processing.construct_router_data_multisample --resume
"""

from __future__ import annotations

import argparse
import hashlib
import json
import logging
import random
import threading
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np
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

# LLM variance tiers based on experiment results.
# High variance LLMs benefit from more samples.
LLM_VARIANCE_TIERS: dict[str, int] = {
    # High variance (>30%): 5 samples
    "Gemini 2.5 Flash Lite": 5,
    "GPT-5.1": 5,
    # Medium variance (20-30%): 3 samples
    "Claude Haiku 4.5": 3,
    "GPT-5 Nano": 3,
    "Gemini 3 Pro Preview": 3,
    "Gemini 3 Flash Preview": 3,
    # Low variance (<20%): 2 samples
    "Claude Opus 4.5": 2,
    # Deterministic (no temperature): 1 sample
    "GPT-5.1 Codex": 1,
}
DEFAULT_N_SAMPLES = 3


def _get_n_samples_for_llm(llm_name: str, default: int = DEFAULT_N_SAMPLES) -> int:
    """Get recommended sample count for an LLM based on variance tier."""
    return LLM_VARIANCE_TIERS.get(llm_name, default)


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 "default"


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


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

    def __init__(
        self,
        initial: int = 10,
        min_concurrent: int = 2,
        max_concurrent: int = 50,
        increase_after: int = 20,
        decrease_factor: float = 0.5,
    ):
        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

    def acquire(self):
        self._semaphore.acquire()

    def release_success(self):
        with self._lock:
            self._success_count += 1
            self._total_successes += 1
            if (
                self._success_count >= self.increase_after
                and self.current < self.max_concurrent
            ):
                self.current += 1
                self._semaphore.release()
                self._semaphore.release()
                self._success_count = 0
                return
        self._semaphore.release()

    def release_failure(self, is_rate_limit: bool = False):
        with self._lock:
            self._total_failures += 1
            self._success_count = 0
            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:
                    self.current = new_limit
                    return
        self._semaphore.release()

    def stats(self) -> dict:
        with self._lock:
            total = self._total_successes + self._total_failures
            return {
                "current_limit": self.current,
                "success_rate": self._total_successes / max(1, total),
            }


@dataclass
class MultiSampleConfig:
    """Configuration for multi-sample data collection."""

    n_samples_default: int = 3
    use_variance_tiers: bool = True  # Use LLM-specific sample counts
    temperature: float = 0.7
    max_tokens: int = 16384
    max_workers: int = 20
    seed: int = 42
    checkpoint_interval: int = 100
    save_interval: int = 20


@dataclass
class AggregatedResult:
    """Aggregated statistics for a (query, LLM) pair across samples."""

    query: str
    llm: str
    task_id: str
    ground_truth: str
    metric: str
    n_samples: int
    effects: list[float]
    effect_mean: float
    effect_std: float
    effect_min: float
    effect_max: float
    responses: list[str]
    latencies: list[float]


class MultiSampleDataBuilder:
    """Build router training data with multi-sample collection."""

    def __init__(
        self,
        qa_path: Path,
        llm_path: Path,
        config: MultiSampleConfig,
        output_dir: Path | None = None,
    ):
        self.qa_data = pd.read_csv(qa_path)
        if self.qa_data.empty:
            raise ValueError("QA dataset is empty")

        with open(llm_path) as f:
            self.llm_description = json.load(f)
        self.llm_names = list(self.llm_description.keys())
        if not self.llm_names:
            raise ValueError("LLM description catalog is empty")

        self.config = config

        # Output directory
        self.output_dir = output_dir or (_MODULE_DIR / "data" / ".multisample_collection")
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Cache and checkpoint files
        self.cache_file = self.output_dir / "sample_cache.json"
        self.checkpoint_file = self.output_dir / "checkpoint.json"
        self.results_file = self.output_dir / "router_data_multisample.csv"

        # Thread-safe cache
        self._cache_lock = threading.Lock()
        self._unsaved_count = 0

        # Load cache
        self.cache = self._load_cache()

        # Rate limiters per provider
        self._rate_limiters: dict[str, AdaptiveRateLimiter] = {
            "openai": AdaptiveRateLimiter(initial=5, max_concurrent=30),
            "google": AdaptiveRateLimiter(initial=10, max_concurrent=50),
            "anthropic": AdaptiveRateLimiter(initial=8, max_concurrent=30),
            "default": AdaptiveRateLimiter(initial=3, max_concurrent=20),
        }

        # LLM engine (lazy loaded)
        self._llm_engine = None

    def _load_cache(self) -> dict:
        if self.cache_file.exists():
            try:
                return json.loads(self.cache_file.read_text(encoding="utf-8"))
            except json.JSONDecodeError:
                logger.warning("Corrupted cache, starting fresh")
                return {}
        return {}

    def _save_cache(self, force: bool = False) -> None:
        with self._cache_lock:
            if force or self._unsaved_count >= self.config.save_interval:
                self.cache_file.write_text(
                    json.dumps(self.cache, indent=2), encoding="utf-8"
                )
                self._unsaved_count = 0

    def _save_checkpoint(self, completed_keys: set[str]) -> None:
        checkpoint = {
            "completed_keys": list(completed_keys),
            "timestamp": time.time(),
            "config": {
                "n_samples_default": self.config.n_samples_default,
                "temperature": self.config.temperature,
            },
        }
        self.checkpoint_file.write_text(json.dumps(checkpoint, indent=2), encoding="utf-8")

    def _load_checkpoint(self) -> set[str]:
        if self.checkpoint_file.exists():
            try:
                checkpoint = json.loads(self.checkpoint_file.read_text())
                return set(checkpoint.get("completed_keys", []))
            except json.JSONDecodeError:
                return set()
        return set()

    def _load_llm_engine(self):
        if self._llm_engine is None:
            from .llm_engine import LLMEngine

            self._llm_engine = LLMEngine(
                llm_names=self.llm_names,
                llm_description=self.llm_description,
            )
        return self._llm_engine

    def _run_single_sample(
        self,
        query: str,
        llm_name: str,
        sample_id: int,
        ground_truth: str,
        metric: str,
    ) -> dict:
        """Run a single sample and return result."""
        cache_key = _query_sample_hash(query, llm_name, sample_id)

        # Check cache
        with self._cache_lock:
            if cache_key in self.cache:
                cached = self.cache[cache_key]
                return {
                    "response": cached["response"],
                    "effect": cached["effect"],
                    "latency": cached["latency"],
                    "from_cache": True,
                }

        # Load engine
        engine = self._load_llm_engine()
        model_id = self.llm_description[llm_name]["model"]
        provider = _get_provider(llm_name)
        limiter = self._rate_limiters.get(provider, self._rate_limiters["default"])

        # Determine temperature (0 for deterministic models)
        temperature = self.config.temperature
        if "codex" in llm_name.lower() or "o1" in model_id.lower():
            temperature = None  # Deterministic models don't support temperature

        # Import prompting function
        from .utils import model_prompting

        # Make API call
        limiter.acquire()
        start = time.perf_counter()
        try:
            response = model_prompting(
                llm_model=model_id,
                prompt=query,
                max_token_num=self.config.max_tokens,
                temperature=temperature,
            )
            limiter.release_success()
        except Exception as e:
            error_msg = str(e).lower()
            is_rate_limit = any(x in error_msg for x in ["rate", "limit", "429", "quota"])
            limiter.release_failure(is_rate_limit)
            logger.warning("API error for %s sample %d: %s", llm_name, sample_id, e)
            response = ""

        latency = time.perf_counter() - start

        # Evaluate response
        effect = engine.eval(
            prediction=response,
            ground_truth=ground_truth,
            metric=metric,
        )

        # Cache result
        result_data = {
            "response": response,
            "effect": effect,
            "latency": latency,
            "timestamp": time.time(),
        }
        with self._cache_lock:
            self.cache[cache_key] = result_data
            self._unsaved_count += 1
        self._save_cache()

        return {
            "response": response,
            "effect": effect,
            "latency": latency,
            "from_cache": False,
        }

    def _collect_samples_for_pair(
        self,
        row_idx: int,
        query: str,
        llm_name: str,
        ground_truth: str,
        metric: str,
        task_id: str,
    ) -> AggregatedResult:
        """Collect multiple samples for a (query, LLM) pair."""
        # Get sample count for this LLM
        if self.config.use_variance_tiers:
            n_samples = _get_n_samples_for_llm(llm_name, self.config.n_samples_default)
        else:
            n_samples = self.config.n_samples_default

        effects = []
        responses = []
        latencies = []

        for sample_id in range(n_samples):
            result = self._run_single_sample(
                query=query,
                llm_name=llm_name,
                sample_id=sample_id,
                ground_truth=ground_truth,
                metric=metric,
            )
            effects.append(result["effect"])
            responses.append(result["response"][:200])  # Truncate for storage
            latencies.append(result["latency"])

        effects_arr = np.array(effects)

        return AggregatedResult(
            query=query,
            llm=llm_name,
            task_id=task_id,
            ground_truth=ground_truth,
            metric=metric,
            n_samples=n_samples,
            effects=effects,
            effect_mean=float(effects_arr.mean()),
            effect_std=float(effects_arr.std()),
            effect_min=float(effects_arr.min()),
            effect_max=float(effects_arr.max()),
            responses=responses,
            latencies=latencies,
        )

    def build(self, resume: bool = True) -> pd.DataFrame:
        """Build the multi-sample router dataset."""
        logger.info("Loading QA data: %d queries", len(self.qa_data))
        logger.info("LLMs: %s", self.llm_names)

        # Build work items: (row_idx, query, llm_name, ground_truth, metric, task_id)
        work_items = []
        for row_idx, row in self.qa_data.iterrows():
            for llm_name in self.llm_names:
                work_items.append(
                    {
                        "row_idx": row_idx,
                        "query": row["query"],
                        "llm_name": llm_name,
                        "ground_truth": row["ground_truth"],
                        "metric": row["metric"],
                        "task_id": row["task_id"],
                    }
                )

        # Load checkpoint for resumption
        completed_keys = self._load_checkpoint() if resume else set()
        if completed_keys:
            logger.info("Resuming with %d completed pairs", len(completed_keys))

        # Filter to uncompleted items
        def pair_key(item):
            return f"{item['row_idx']}||{item['llm_name']}"

        work_items = [item for item in work_items if pair_key(item) not in completed_keys]
        logger.info("Pairs to process: %d", len(work_items))

        if not work_items:
            logger.info("All pairs completed. Loading from cache...")
            return self._build_dataframe_from_cache()

        # Shuffle for better provider distribution
        random.seed(self.config.seed)
        random.shuffle(work_items)

        # Process with progress bar
        results: list[AggregatedResult] = []
        checkpoint_count = 0

        with ThreadPoolExecutor(max_workers=self.config.max_workers) as executor:
            futures = {}
            for item in work_items:
                future = executor.submit(
                    self._collect_samples_for_pair,
                    item["row_idx"],
                    item["query"],
                    item["llm_name"],
                    item["ground_truth"],
                    item["metric"],
                    item["task_id"],
                )
                futures[future] = item

            with tqdm(total=len(work_items), desc="Collecting") as pbar:
                for future in as_completed(futures):
                    item = futures[future]
                    try:
                        result = future.result()
                        results.append(result)
                        completed_keys.add(pair_key(item))
                        checkpoint_count += 1

                        if checkpoint_count >= self.config.checkpoint_interval:
                            self._save_checkpoint(completed_keys)
                            self._save_cache(force=True)
                            checkpoint_count = 0

                    except Exception as e:
                        logger.warning(
                            "Failed: row %d, %s: %s",
                            item["row_idx"],
                            item["llm_name"],
                            e,
                        )
                    pbar.update(1)

        # Final save
        self._save_cache(force=True)
        self._save_checkpoint(completed_keys)

        # Build and save dataframe
        df = self._build_dataframe(results)
        df.to_csv(self.results_file, index=False)
        logger.info("Saved %d records to %s", len(df), self.results_file)

        # Clean up checkpoint
        if self.checkpoint_file.exists():
            self.checkpoint_file.unlink()

        return df

    def _build_dataframe(self, results: list[AggregatedResult]) -> pd.DataFrame:
        """Build dataframe from aggregated results."""
        records = []
        for r in results:
            records.append(
                {
                    "task_id": r.task_id,
                    "query": r.query,
                    "ground_truth": r.ground_truth,
                    "metric": r.metric,
                    "llm": r.llm,
                    "n_samples": r.n_samples,
                    "effect": r.effect_mean,  # Use mean as primary effect
                    "effect_mean": r.effect_mean,
                    "effect_std": r.effect_std,
                    "effect_min": r.effect_min,
                    "effect_max": r.effect_max,
                    "is_high_variance": r.effect_std > 0.1,
                    "latency_mean": np.mean(r.latencies),
                }
            )
        return pd.DataFrame(records)

    def _build_dataframe_from_cache(self) -> pd.DataFrame:
        """Build dataframe from existing cache."""
        # Group cache entries by (query, llm) pair
        pair_effects: dict[tuple, list] = defaultdict(list)
        pair_meta: dict[tuple, dict] = {}

        # This requires reconstructing from cache - simplified version
        logger.warning("Cache-only reconstruction not fully implemented")
        return pd.DataFrame()


def main():
    parser = argparse.ArgumentParser(
        description="Multi-sample router data collection"
    )
    parser.add_argument(
        "--qa-path",
        type=Path,
        default=_MODULE_DIR / "data" / "unified_qa_data.csv",
        help="Path to unified QA data",
    )
    parser.add_argument(
        "--llm-path",
        type=Path,
        default=_MODULE_DIR / "configs" / "LLM_Descriptions.json",
        help="Path to LLM descriptions",
    )
    parser.add_argument(
        "--n-samples",
        type=int,
        default=3,
        help="Default samples per pair (overridden by variance tiers)",
    )
    parser.add_argument(
        "--no-variance-tiers",
        action="store_true",
        help="Disable variance-based sample counts",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="Sampling temperature",
    )
    parser.add_argument(
        "--max-workers",
        type=int,
        default=20,
        help="Max concurrent API calls",
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Resume from checkpoint",
    )
    parser.add_argument(
        "--no-resume",
        action="store_true",
        help="Start fresh (ignore checkpoint)",
    )

    args = parser.parse_args()

    config = MultiSampleConfig(
        n_samples_default=args.n_samples,
        use_variance_tiers=not args.no_variance_tiers,
        temperature=args.temperature,
        max_workers=args.max_workers,
    )

    builder = MultiSampleDataBuilder(
        qa_path=args.qa_path,
        llm_path=args.llm_path,
        config=config,
    )

    resume = not args.no_resume
    df = builder.build(resume=resume)

    # Print summary
    print("\n" + "=" * 60)
    print("MULTI-SAMPLE DATA COLLECTION COMPLETE")
    print("=" * 60)
    print(f"Total records: {len(df)}")
    print(f"Unique queries: {df['query'].nunique()}")
    print(f"LLMs: {df['llm'].nunique()}")
    print(f"High variance pairs: {df['is_high_variance'].sum()} ({100*df['is_high_variance'].mean():.1f}%)")
    print(f"Mean effect: {df['effect'].mean():.4f}")
    print(f"Output: {builder.results_file}")
    print("=" * 60)


if __name__ == "__main__":
    main()
