from pathlib import Path

import hashlib
import json
import logging
import random
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
import threading
import pandas as pd
from tqdm import tqdm
from transformers import AutoTokenizer
import yaml

# Default max concurrent API calls across all threads.
# Per-provider AIMD rate limiters handle backpressure, so this global cap
# can be set high.  The limiters start conservatively and probe upward.
DEFAULT_MAX_CONCURRENT_API_CALLS = 80

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

# Resolve module directory for config path resolution
_MODULE_DIR = Path(__file__).resolve().parent.parent


@dataclass
class ProgressTracker:
    """Track progress, throughput, and per-LLM statistics."""

    total_queries: int
    llm_names: list[str]
    start_time: float = field(default_factory=time.time)

    # Counters (thread-safe via locks)
    _completed: int = 0
    _cache_hits: int = 0
    _api_calls: int = 0
    _total_cost: float = 0.0
    _llm_success: dict = field(default_factory=lambda: defaultdict(int))
    _llm_failures: dict = field(default_factory=lambda: defaultdict(int))
    _lock: threading.Lock = field(default_factory=threading.Lock)

    def record_query(
        self,
        query_records: list[dict],
        cache_hits: int,
        api_calls: int,
    ) -> None:
        """Record completion of a query across all LLMs."""
        with self._lock:
            self._completed += 1
            self._cache_hits += cache_hits
            self._api_calls += api_calls
            for rec in query_records:
                self._total_cost += rec.get("cost", 0.0)

    def record_llm_result(self, llm_name: str, success: bool) -> None:
        """Record success/failure for a specific LLM."""
        with self._lock:
            if success:
                self._llm_success[llm_name] += 1
            else:
                self._llm_failures[llm_name] += 1

    def get_stats(self) -> dict:
        """Return current statistics snapshot."""
        with self._lock:
            elapsed = time.time() - self.start_time
            completed = self._completed
            remaining = self.total_queries - completed

            # Throughput and ETA
            qps = completed / elapsed if elapsed > 0 else 0.0
            eta_sec = remaining / qps if qps > 0 else 0.0

            return {
                "completed": completed,
                "total": self.total_queries,
                "pct": 100 * completed / self.total_queries if self.total_queries else 0,
                "qps": qps,
                "eta_min": eta_sec / 60,
                "cache_hits": self._cache_hits,
                "api_calls": self._api_calls,
                "cache_rate": self._cache_hits / (self._cache_hits + self._api_calls)
                if (self._cache_hits + self._api_calls) > 0
                else 0.0,
                "total_cost": self._total_cost,
                "elapsed_min": elapsed / 60,
            }

    def format_postfix(self) -> dict:
        """Return formatted dict for tqdm postfix."""
        stats = self.get_stats()
        return {
            "qps": f"{stats['qps']:.1f}",
            "eta": f"{stats['eta_min']:.0f}m",
            "cache": f"{stats['cache_rate']*100:.0f}%",
            "cost": f"${stats['total_cost']:.2f}",
        }

    def log_summary(self) -> None:
        """Log final summary statistics."""
        stats = self.get_stats()
        logger.info("=" * 60)
        logger.info("PROCESSING COMPLETE")
        logger.info("=" * 60)
        logger.info(
            "Queries: %d/%d (%.1f%%)",
            stats["completed"],
            stats["total"],
            stats["pct"],
        )
        logger.info(
            "Time: %.1f min | Throughput: %.2f queries/sec",
            stats["elapsed_min"],
            stats["qps"],
        )
        logger.info(
            "API calls: %d | Cache hits: %d (%.0f%% cache rate)",
            stats["api_calls"],
            stats["cache_hits"],
            stats["cache_rate"] * 100,
        )
        logger.info("Estimated cost: $%.2f", stats["total_cost"])

        # Per-LLM breakdown
        logger.info("-" * 40)
        logger.info("Per-LLM Statistics:")
        with self._lock:
            for llm in self.llm_names:
                success = self._llm_success.get(llm, 0)
                fail = self._llm_failures.get(llm, 0)
                total = success + fail
                rate = success / total * 100 if total > 0 else 0
                status = "✓" if fail == 0 else "!"
                logger.info(
                    "  %s %-20s: %4d success, %3d fail (%.0f%%)",
                    status,
                    llm[:20],
                    success,
                    fail,
                    rate,
                )


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,
        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):
        """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("Rate limit increased to %d", 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:
                    slots_to_drain = self.current - new_limit
                    self.current = new_limit
                    logger.warning(
                        "Rate limit decreased to %d (draining %d slots)", self.current, slots_to_drain
                    )
                    return  # Don't release - let semaphore drain naturally

        self._semaphore.release()

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


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"


# Handle both package import and direct script execution
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 _text_hash(text: str) -> str:
    """Generate a unique hash for a text (query or task description)."""
    return hashlib.sha256(text.encode()).hexdigest()[:16]


class data_building:
    def __init__(
        self,
        qa_path,
        llm_path,
        config,
        cache_dir: Path | None = None,
        max_query_workers: int = 10,
        max_llm_workers: int = 7,
        max_concurrent_api_calls: int = DEFAULT_MAX_CONCURRENT_API_CALLS,
        skip_embeddings: bool = False,
    ):
        self.qa_data = pd.read_csv(qa_path)
        if self.qa_data.empty:
            raise ValueError("Unified QA dataset is empty; run multidata_unify first.")

        self.llm_description = loadjson(llm_path)
        self.llm_names = list(self.llm_description.keys())
        if not self.llm_names:
            raise ValueError("LLM description catalog is empty.")

        self.all_llm_description = [
            self.llm_description[name]["feature"] for name in self.llm_names
        ]
        self.MyLLMEngine = LLMEngine(
            llm_names=self.llm_names, llm_description=self.llm_description
        )
        self.config = config
        self.max_query_workers = max_query_workers
        self.max_llm_workers = max_llm_workers

        # Setup cache directory
        self.cache_dir = cache_dir or (_MODULE_DIR / "data" / ".router_cache")
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.response_cache_file = self.cache_dir / "response_cache.json"
        self.checkpoint_file = self.cache_dir / "checkpoint.json"

        # Thread-safe cache access
        self._cache_lock = threading.Lock()
        self._records_lock = threading.Lock()

        # Adaptive rate limiters per provider using AIMD congestion control.
        # Initial limits are conservative; the algorithm will probe upward on success
        # and back off quickly on 429s.
        self._rate_limiters: dict[str, AdaptiveRateLimiter] = {
            "openai": AdaptiveRateLimiter(initial=20, min_concurrent=2, max_concurrent=50),
            "google": AdaptiveRateLimiter(initial=20, min_concurrent=2, max_concurrent=100),
            "anthropic": AdaptiveRateLimiter(initial=15, min_concurrent=2, max_concurrent=50),
            "default": AdaptiveRateLimiter(initial=5, min_concurrent=1, max_concurrent=20),
        }
        self._max_concurrent_api_calls = max_concurrent_api_calls
        self._skip_embeddings = skip_embeddings

        # Track in-flight calls per provider for visibility
        self._in_flight = defaultdict(int)
        self._in_flight_lock = threading.Lock()

        # Incremental cache saving: save after every N new API responses
        self._unsaved_responses = 0
        self._save_every_n_responses = 5  # Save cache to disk every 5 new responses
        self._last_cache_save = time.time()
        self._max_save_interval_sec = 60  # Force save at least every 60 seconds

        # Load existing caches
        self.response_cache = self._load_cache()
        self.embedding_cache = self._load_embedding_cache()

        # Progress tracker (initialized in construct_data_with_LLM)
        self.tracker: ProgressTracker | None = None

        # Tokenizer for multi_news truncation
        self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
        self.max_tokens = int(getattr(self.tokenizer, "model_max_length", 512) or 512)

    def _load_cache(self) -> dict:
        """Load response cache from disk."""
        if self.response_cache_file.exists():
            try:
                return json.loads(self.response_cache_file.read_text(encoding="utf-8"))
            except json.JSONDecodeError:
                return {}
        return {}

    def _load_embedding_cache(self) -> dict:
        """Load embedding cache from disk."""
        embedding_cache_file = self.cache_dir / "embedding_cache.json"
        if embedding_cache_file.exists():
            try:
                return json.loads(embedding_cache_file.read_text(encoding="utf-8"))
            except json.JSONDecodeError:
                return {}
        return {}

    def _save_cache(self):
        """Save response cache to disk."""
        with self._cache_lock:
            self.response_cache_file.write_text(
                json.dumps(self.response_cache, indent=2), encoding="utf-8"
            )
            self._unsaved_responses = 0
            self._last_cache_save = time.time()

    def _maybe_save_cache(self):
        """Save cache if enough new responses accumulated or time threshold passed."""
        time_since_save = time.time() - self._last_cache_save
        should_save = (
            self._unsaved_responses >= self._save_every_n_responses
            or time_since_save >= self._max_save_interval_sec
        )
        if should_save:
            self._save_cache()

    def _save_checkpoint(self, processed_indices: set, records: list):
        """Save checkpoint for resumption."""
        checkpoint = {
            "processed_indices": sorted(processed_indices),
            "records_count": len(records),
        }
        self.checkpoint_file.write_text(json.dumps(checkpoint), encoding="utf-8")

        # Also save partial records
        with self._records_lock:
            partial_df = pd.DataFrame.from_records(records)
        partial_path = self.cache_dir / "partial_records.csv"
        partial_df.to_csv(partial_path, index=False)

    def _load_checkpoint(self) -> tuple[set, list]:
        """Load checkpoint if exists.

        IMPORTANT: processed_indices is derived from actual records, not the
        checkpoint file. This ensures consistency: a query is only considered
        "processed" if it has complete records for ALL LLMs. The checkpoint
        file's processed_indices is ignored to avoid race condition bugs where
        indices were marked processed before records were persisted.
        """
        partial_path = self.cache_dir / "partial_records.csv"

        if not partial_path.exists():
            return set(), []

        partial_df = pd.read_csv(partial_path, low_memory=False)
        if partial_df.empty:
            return set(), []

        # Derive processed_indices from records with COMPLETE LLM coverage.
        # A query is only "done" if it has exactly len(llm_names) records.
        llm_count = len(self.llm_names)
        records_per_query = partial_df.groupby("row_idx").size()
        complete_queries = records_per_query[records_per_query == llm_count].index
        processed = set(complete_queries.tolist())

        # Filter to only keep records from complete queries to avoid duplicates
        # when incomplete queries are re-processed.
        complete_df = partial_df[partial_df["row_idx"].isin(processed)]
        records = complete_df.to_dict("records")

        # Log discrepancy if checkpoint exists (for debugging)
        if self.checkpoint_file.exists():
            checkpoint = json.loads(self.checkpoint_file.read_text(encoding="utf-8"))
            checkpoint_processed = set(checkpoint.get("processed_indices", []))
            if checkpoint_processed != processed:
                missing = checkpoint_processed - processed
                logger.warning(
                    "Checkpoint had %d indices marked processed, but only %d have "
                    "complete records. Recovering %d queries from cache.",
                    len(checkpoint_processed),
                    len(processed),
                    len(missing),
                )

        logger.info(
            "Resuming with %d complete queries and %d cached records",
            len(processed),
            len(records),
        )
        return processed, records

    def _query_single_llm(self, query: str, llm_idx: int, llm_name: str) -> dict:
        """Query a single LLM and return result dict. Thread-safe."""
        cache_key = _query_hash(query, llm_name)

        # Check cache first (thread-safe read)
        # Skip empty cached responses - these are failed API calls that need retry
        with self._cache_lock:
            if cache_key in self.response_cache:
                cached = self.response_cache[cache_key]
                cached_response = cached.get("response", "")
                if cached_response and cached_response.strip():
                    if self.tracker:
                        self.tracker.record_llm_result(llm_name, success=True)
                    return {
                        "llm_idx": llm_idx,
                        "llm_name": llm_name,
                        "response": cached_response,
                        "latency": cached["latency"],
                        "from_cache": True,
                    }
                # Empty cached response - remove it and re-query
                del self.response_cache[cache_key]

        # Make API call with adaptive per-provider concurrency control
        provider = _get_provider(llm_name)
        limiter = self._rate_limiters.get(provider, self._rate_limiters["default"])

        limiter.acquire()
        is_rate_limit = False
        try:
            with self._in_flight_lock:
                self._in_flight[provider] += 1
                in_flight_str = " ".join(f"{k}:{v}" for k, v in self._in_flight.items() if v > 0)
            try:
                response, latency = self.MyLLMEngine.get_llm_response_with_latency(
                    query=query,
                    llm_idx=llm_idx,
                )
                success = bool(response)  # Non-empty response = success
                logger.info("✓ %s in %.1fs [%s]", llm_name, latency, in_flight_str)
            except Exception as e:
                error_str = str(e).lower()
                is_rate_limit = "429" in error_str or "rate" in error_str
                logger.warning("✗ %s failed: %s", llm_name, e)
                response, latency, success = "", 0.0, False
            finally:
                with self._in_flight_lock:
                    self._in_flight[provider] -= 1
        finally:
            # Adaptive concurrency: adjust limits based on success/failure
            if success:
                limiter.release_success()
            else:
                limiter.release_failure(is_rate_limit=is_rate_limit)

        if self.tracker:
            self.tracker.record_llm_result(llm_name, success=success)

        # Update cache (thread-safe write) - only cache successful responses
        if success:
            with self._cache_lock:
                self.response_cache[cache_key] = {
                    "response": response,
                    "latency": latency,
                }
                self._unsaved_responses += 1
            # Persist to disk frequently to avoid losing progress
            self._maybe_save_cache()

        return {
            "llm_idx": llm_idx,
            "llm_name": llm_name,
            "response": response,
            "latency": latency,
            "from_cache": False,
        }

    def _process_single_query(self, row_idx: int, row) -> list[dict]:
        """Process a single query across all LLMs. Returns list of records."""
        task_id = row.task_id
        query = row.query
        task_description = row.task_description

        # Truncate multi_news queries
        if task_id == "multi_news":
            tokens = self.tokenizer.tokenize(query)
            query = self.tokenizer.convert_tokens_to_string(tokens[: self.max_tokens])

        # Get embeddings from cache if available, compute only if needed
        query_hash = _text_hash(query)
        task_hash = _text_hash(task_description)

        # Try to get embeddings from cache first (thread-safe read)
        query_embedding_str = self.embedding_cache.get(query_hash)
        task_embedding_str = self.embedding_cache.get(task_hash)

        if self._skip_embeddings:
            # Use cached embeddings if available, otherwise leave empty
            query_embedding = query_embedding_str
            task_description_embedding = task_embedding_str
        else:
            # Compute missing embeddings
            if query_embedding_str is None:
                query_embedding_str = str(get_embedding([query]).tolist())
            if task_embedding_str is None:
                task_embedding_str = str(get_embedding([task_description]).tolist())
            query_embedding = query_embedding_str
            task_description_embedding = task_embedding_str
        ground_truth = row.ground_truth
        metric = row.metric

        # Query all LLMs in parallel
        llm_results = []
        with ThreadPoolExecutor(max_workers=self.max_llm_workers) as executor:
            futures = {
                executor.submit(
                    self._query_single_llm, query, llm_idx, llm_name
                ): llm_name
                for llm_idx, llm_name in enumerate(self.llm_names)
            }

            for future in as_completed(futures):
                try:
                    result = future.result()
                    llm_results.append(result)
                except Exception as e:
                    llm_name = futures[future]
                    print(f"\nError querying {llm_name}: {e}")
                    llm_results.append(
                        {
                            "llm_idx": self.llm_names.index(llm_name),
                            "llm_name": llm_name,
                            "response": "",
                            "latency": 0.0,
                            "from_cache": False,
                        }
                    )

        # Sort by llm_idx for consistent ordering
        llm_results.sort(key=lambda x: x["llm_idx"])

        # Build records for this query
        query_records = []
        for result in llm_results:
            llm_idx = result["llm_idx"]
            llm_name = result["llm_name"]
            response = result["response"]
            latency = result["latency"]
            from_cache = result.get("from_cache", False)

            reward = self.MyLLMEngine.eval(
                prediction=response, ground_truth=ground_truth, metric=metric
            )
            cost = self.MyLLMEngine.compute_cost(llm_idx, query, response)

            query_records.append(
                {
                    "row_idx": row_idx,
                    "task_id": task_id,
                    "task_description": task_description,
                    "task_description_embedding": task_description_embedding or "",
                    "query": query,
                    "query_embedding": query_embedding or "",
                    "ground_truth": ground_truth,
                    "metric": metric,
                    "llm": llm_name,
                    "effect": reward,
                    "cost": cost,
                    "latency": latency,
                    "from_cache": from_cache,  # Track for progress reporting
                }
            )

        return query_records

    def _get_cached_response(self, query: str, llm_name: str) -> dict | None:
        """Check if a response is cached. Returns cached data or None."""
        cache_key = _query_hash(query, llm_name)
        with self._cache_lock:
            if cache_key in self.response_cache:
                cached = self.response_cache[cache_key]
                cached_response = cached.get("response", "")
                if cached_response and cached_response.strip():
                    return cached
        return None

    def _process_single_llm_call(
        self, row_idx: int, row, llm_idx: int, llm_name: str
    ) -> dict:
        """Process a single (query, LLM) pair. Fully independent, no blocking."""
        task_id = row.task_id
        query = row.query
        task_description = row.task_description

        # Truncate multi_news queries
        if task_id == "multi_news":
            tokens = self.tokenizer.tokenize(query)
            query = self.tokenizer.convert_tokens_to_string(tokens[: self.max_tokens])

        # Get embeddings from cache if available
        query_hash = _text_hash(query)
        task_hash = _text_hash(task_description)
        query_embedding = self.embedding_cache.get(query_hash)
        task_embedding = self.embedding_cache.get(task_hash)

        ground_truth = row.ground_truth
        metric = row.metric

        # Query this single LLM
        result = self._query_single_llm(query, llm_idx, llm_name)

        # Evaluate the response
        reward = self.MyLLMEngine.eval(
            prediction=result["response"],
            ground_truth=ground_truth,
            metric=metric,
        )

        return {
            "row_idx": row_idx,
            "task_id": task_id,
            "task_description": task_description,
            "task_description_embedding": task_embedding or "",
            "query": query,
            "query_embedding": query_embedding or "",
            "ground_truth": ground_truth,
            "metric": metric,
            "llm": llm_name,
            "llm_idx": llm_idx,
            "effect": reward,
            "cost": self.MyLLMEngine.compute_cost(llm_idx, query, result["response"]),
            "latency": result["latency"],
            "from_cache": result["from_cache"],
        }

    def construct_data_with_LLM(self):
        """Process all (query, LLM) pairs independently - no blocking between providers."""
        total_rows = len(self.qa_data)
        total_responses = total_rows * len(self.llm_names)

        # Build flat list of all (row, llm) pairs that need processing
        # Check cache to skip already-completed pairs
        pairs_to_process = []
        cached_count = 0

        for row_idx, row in enumerate(self.qa_data.itertuples()):
            query = row.query
            # Truncate multi_news for cache check
            if row.task_id == "multi_news":
                tokens = self.tokenizer.tokenize(query)
                query = self.tokenizer.convert_tokens_to_string(tokens[: self.max_tokens])

            for llm_idx, llm_name in enumerate(self.llm_names):
                if self._get_cached_response(query, llm_name) is None:
                    pairs_to_process.append((row_idx, row, llm_idx, llm_name))
                else:
                    cached_count += 1

        logger.info(
            "Total: %d responses (%d cached, %d to process)",
            total_responses,
            cached_count,
            len(pairs_to_process),
        )

        if not pairs_to_process:
            logger.info("All responses cached. Assembling final data...")
        else:
            # Hybrid stratified shuffle: interleave tasks AND prioritize short queries.
            #
            # Why hybrid?
            # - Pure stratified: accurate progress, but 203x query length variance
            #   means ETA is still misleading (short tasks fast, long tasks slow)
            # - Pure short-first: maximum throughput, but progress bar lies
            # - Hybrid: interleave tasks, but within each "round" process shortest first
            #
            # Algorithm:
            # 1. Group pairs by task_id
            # 2. Sort within each task by query length (shortest first)
            # 3. Round-robin across tasks, but within each round sort by query length
            logger.info("Hybrid stratified shuffle (task interleave + short-first)...")

            # Group by task_id
            by_task: dict[str, list] = defaultdict(list)
            for pair in pairs_to_process:
                task_id = pair[1].task_id  # row.task_id
                by_task[task_id].append(pair)

            # Sort within each task by query length (shortest first for faster processing)
            for task_pairs in by_task.values():
                task_pairs.sort(key=lambda p: len(p[1].query))

            # Round-robin interleave across tasks
            # Within each round, collect one query from each task, then sort by length
            seed = self.config.get("seed", 1881)
            rng = random.Random(seed)
            task_iterators = {k: iter(v) for k, v in by_task.items()}
            task_names = list(by_task.keys())
            rng.shuffle(task_names)  # Randomize task order for fairness

            pairs_to_process = []
            round_num = 0
            while task_iterators:
                # Collect one query from each remaining task for this round
                round_pairs = []
                empty_tasks = []
                for task in task_names:
                    if task not in task_iterators:
                        continue
                    try:
                        round_pairs.append(next(task_iterators[task]))
                    except StopIteration:
                        empty_tasks.append(task)
                for task in empty_tasks:
                    del task_iterators[task]
                    task_names.remove(task)

                # Sort this round's pairs by query length (shortest first)
                round_pairs.sort(key=lambda p: len(p[1].query))
                pairs_to_process.extend(round_pairs)
                round_num += 1

            logger.info(
                "Shuffled %d pairs across %d tasks in %d rounds (seed=%d)",
                len(pairs_to_process),
                len(by_task),
                round_num,
                seed,
            )

            # Group by provider for logging
            by_provider = defaultdict(int)
            by_llm = defaultdict(int)
            for _, _, _, llm_name in pairs_to_process:
                by_llm[llm_name] += 1
                by_provider[_get_provider(llm_name)] += 1
            logger.info("Pending by provider: %s", dict(by_provider))
            logger.info("Pending by LLM: %s", dict(by_llm))

            # Process all pairs in parallel - flat structure, no nesting
            # Each pair is independent - slowness in one LLM doesn't block others
            max_workers = self._max_concurrent_api_calls
            logger.info("Processing with %d parallel workers", max_workers)

            pbar = tqdm(total=len(pairs_to_process), desc="API calls")
            completed_count = 0
            save_interval = 50  # Save every N responses

            # Heartbeat thread to show activity while waiting for slow models
            heartbeat_stop = threading.Event()
            def heartbeat():
                while not heartbeat_stop.wait(10):  # Every 10 seconds
                    with self._in_flight_lock:
                        in_flight = dict(self._in_flight)
                    total_in_flight = sum(in_flight.values())
                    if total_in_flight > 0:
                        status = " ".join(f"{k}:{v}" for k, v in in_flight.items() if v > 0)
                        logger.info("⏳ Waiting... [%s] %d in-flight", status, total_in_flight)

            heartbeat_thread = threading.Thread(target=heartbeat, daemon=True)
            heartbeat_thread.start()

            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                futures = {
                    executor.submit(
                        self._process_single_llm_call, row_idx, row, llm_idx, llm_name
                    ): (row_idx, llm_name)
                    for row_idx, row, llm_idx, llm_name in pairs_to_process
                }

                for future in as_completed(futures):
                    row_idx, llm_name = futures[future]
                    try:
                        result = future.result()
                        # Result is auto-cached in _query_single_llm
                        completed_count += 1
                        pbar.update(1)

                        # Periodic save
                        if completed_count % save_interval == 0:
                            self._save_cache()
                            logger.info(
                                "Progress: %d/%d responses (%.1f%%)",
                                completed_count,
                                len(pairs_to_process),
                                100 * completed_count / len(pairs_to_process),
                            )

                    except Exception as e:
                        logger.error(
                            "Error processing row %d, %s: %s", row_idx, llm_name, e
                        )
                        pbar.update(1)

            pbar.close()
            heartbeat_stop.set()  # Stop heartbeat thread
            self._save_cache()

            # Log adaptive rate limiter statistics
            logger.info("-" * 40)
            logger.info("Adaptive Rate Limiter Statistics:")
            for provider, limiter in self._rate_limiters.items():
                stats = limiter.stats()
                if stats["total_successes"] + stats["total_failures"] > 0:
                    logger.info(
                        "  %-10s: limit=%d, success=%d, fail=%d (%.0f%% success rate)",
                        provider,
                        stats["current_limit"],
                        stats["total_successes"],
                        stats["total_failures"],
                        stats["success_rate"] * 100,
                    )

        # Assemble final records from cache
        logger.info("Assembling final records from cache...")
        records = []
        for row_idx, row in enumerate(tqdm(
            self.qa_data.itertuples(), total=total_rows, desc="Assembling"
        )):
            query = row.query
            task_id = row.task_id
            task_description = row.task_description
            ground_truth = row.ground_truth
            metric = row.metric

            if task_id == "multi_news":
                tokens = self.tokenizer.tokenize(query)
                query = self.tokenizer.convert_tokens_to_string(tokens[: self.max_tokens])

            # Get embeddings
            query_hash = _text_hash(query)
            task_hash = _text_hash(task_description)
            query_embedding = self.embedding_cache.get(query_hash, "")
            task_embedding = self.embedding_cache.get(task_hash, "")

            for llm_idx, llm_name in enumerate(self.llm_names):
                cached = self._get_cached_response(query, llm_name)
                if cached:
                    response = cached["response"]
                    latency = cached["latency"]
                else:
                    response = ""
                    latency = 0.0
                    logger.warning(
                        "Missing cache for row %d, %s", row_idx, llm_name
                    )

                reward = self.MyLLMEngine.eval(
                    prediction=response, ground_truth=ground_truth, metric=metric
                )

                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": reward,
                    "cost": self.MyLLMEngine.compute_cost(
                        llm_idx, query, response
                    ),
                    "latency": latency,
                })

        logger.info("Assembled %d records", len(records))

        df = pd.DataFrame.from_records(records)

        if df.empty:
            raise ValueError("Router dataset construction produced no rows.")

        if len(df) % len(self.llm_names) != 0:
            raise ValueError(
                "Router dataset is malformed: total rows are not divisible by LLM count."
            )

        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:
            # Clip latency outliers per (task_id, llm) pair to preserve cross-model differences.
            # This ensures slow models (e.g., Claude Opus) don't have their natural latency
            # clipped against fast models (e.g., Gemini Flash) within the same task.
            df["latency"] = df.groupby(["task_id", "llm"])["latency"].transform(
                lambda s: s.clip(upper=s.quantile(0.99))
            )
            # Normalize per task to make latency comparable across tasks
            df["latency"] = df.groupby("task_id")["latency"].transform(_minmax)

        output_path = Path(self.config["saved_router_data_path"])
        output_path.parent.mkdir(parents=True, exist_ok=True)
        df.to_csv(output_path, index=False)

        # Skip LLM description embedding when deferring embeddings
        if not self._skip_embeddings:
            llm_description_embedding = get_embedding(self.all_llm_description)
            savepkl(llm_description_embedding, self.config["llm_embedding_path"])

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

        router_metadata = {
            "llm_names": self.llm_names,
            "total_rows": int(len(df)),
            "total_queries": int(len(df) // len(self.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(router_metadata, indent=2), encoding="utf-8")

        # Clean up checkpoint after successful completion
        if self.checkpoint_file.exists():
            self.checkpoint_file.unlink()
        partial_path = self.cache_dir / "partial_records.csv"
        if partial_path.exists():
            partial_path.unlink()

        logger.info("Output saved: %s", output_path)
        logger.info(
            "Total records: %d | Total queries: %d",
            len(df),
            len(df) // len(self.llm_names),
        )


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

    # Add data_processing to path for direct execution
    sys.path.insert(0, str(Path(__file__).resolve().parent))

    parser = argparse.ArgumentParser(description="Build router training data")
    parser.add_argument(
        "--config",
        default="config_v2_2026.yaml",
        help="Config filename in configs/ dir (default: config_v2_2026.yaml)",
    )
    parser.add_argument(
        "--max-query-workers",
        type=int,
        default=10,
        help="Max parallel queries to process (default: 10)",
    )
    parser.add_argument(
        "--max-llm-workers",
        type=int,
        default=7,
        help="Max parallel LLM calls per query (default: 7)",
    )
    parser.add_argument(
        "--skip-embeddings",
        action="store_true",
        help="Skip embedding computation (use when embeddings are pre-generated)",
    )
    args = parser.parse_args()

    os.environ["KMP_DUPLICATE_LIB_OK"] = "True"

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

    # Resolve relative paths in config to absolute 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])

    builder = data_building(
        qa_path=config["unified_qa_data_path"],
        llm_path=config["llm_description_path"],
        config=config,
        max_query_workers=args.max_query_workers,
        max_llm_workers=args.max_llm_workers,
        skip_embeddings=args.skip_embeddings,
    )
    builder.construct_data_with_LLM()
