#!/usr/bin/env python3
"""Compute embeddings from multiple sources with parallelism and resumption.

Features:
- Multiple embedding sources: Gemini, OpenAI, local SentenceTransformers
- Matryoshka dimension support (256, 512, 768, 1024, 3072)
- Parallel batch processing with adaptive concurrency
- Automatic resumption via checkpointed cache
- Large batch sizes for throughput

Usage:
    cd /Users/jq/Downloads/takumi/eb1/eb1
    PYTHONPATH=embergraphrouter uv run python -m embergraphrouter.data_processing.compute_multi_embeddings

    # All sources with parallelism:
    uv run python -m embergraphrouter.data_processing.compute_multi_embeddings --sources all --parallel 4

    # Resume interrupted run:
    uv run python -m embergraphrouter.data_processing.compute_multi_embeddings --resume
"""

import argparse
import hashlib
import json
import logging
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from threading import Lock
from typing import Dict, List, Optional, Tuple

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

# Cache file lock for thread-safe writes
_cache_lock = Lock()

# Embedding source configurations with Matryoshka dimensions
# All models tested and confirmed working with Ember
EMBEDDING_SOURCES = {
    # ============ Remote APIs (via Ember) ============
    # Gemini - Google's embedding model (native 768, Matryoshka down to 1)
    # Aligned dimensions: 256, 512, 768 for cross-model ablation
    # Max batch: 2048 texts per request
    "gemini_256": {
        "model_id": "google/gemini-embedding-001",
        "dimension": 256,
        "column_suffix": "gemini_256",
        "type": "remote",
        "batch_size": 2048,
    },
    "gemini_512": {
        "model_id": "google/gemini-embedding-001",
        "dimension": 512,
        "column_suffix": "gemini_512",
        "type": "remote",
        "batch_size": 2048,
    },
    "gemini_768": {
        "model_id": "google/gemini-embedding-001",
        "dimension": 768,
        "column_suffix": "gemini_768",
        "type": "remote",
        "batch_size": 1024,
    },
    # OpenAI text-embedding-3-large (native 3072, Matryoshka down to 1)
    # Aligned dimensions: 256, 512, 768 for cross-model ablation + 1536, 3072 for scaling
    # Adaptive batch sizing will auto-adjust (starts aggressive, falls back if needed)
    "openai_large_256": {
        "model_id": "openai/text-embedding-3-large",
        "dimension": 256,
        "column_suffix": "openai_large_256",
        "type": "remote",
        "batch_size": 512,
    },
    "openai_large_512": {
        "model_id": "openai/text-embedding-3-large",
        "dimension": 512,
        "column_suffix": "openai_large_512",
        "type": "remote",
        "batch_size": 512,
    },
    "openai_large_768": {
        "model_id": "openai/text-embedding-3-large",
        "dimension": 768,
        "column_suffix": "openai_large_768",
        "type": "remote",
        "batch_size": 512,
    },
    "openai_large_1536": {
        "model_id": "openai/text-embedding-3-large",
        "dimension": 1536,
        "column_suffix": "openai_large_1536",
        "type": "remote",
        "batch_size": 512,
    },
    "openai_large_3072": {
        "model_id": "openai/text-embedding-3-large",
        "dimension": 3072,
        "column_suffix": "openai_large_3072",
        "type": "remote",
        "batch_size": 512,
    },
    # OpenAI text-embedding-3-small (native 1536, Matryoshka down to 1)
    # Aligned: 256, 512, 768 for cross-model + 1536 native
    "openai_small_256": {
        "model_id": "openai/text-embedding-3-small",
        "dimension": 256,
        "column_suffix": "openai_small_256",
        "type": "remote",
        "batch_size": 1024,
    },
    "openai_small_512": {
        "model_id": "openai/text-embedding-3-small",
        "dimension": 512,
        "column_suffix": "openai_small_512",
        "type": "remote",
        "batch_size": 1024,
    },
    "openai_small_768": {
        "model_id": "openai/text-embedding-3-small",
        "dimension": 768,
        "column_suffix": "openai_small_768",
        "type": "remote",
        "batch_size": 1024,
    },
    "openai_small_1536": {
        "model_id": "openai/text-embedding-3-small",
        "dimension": 1536,
        "column_suffix": "openai_small_1536",
        "type": "remote",
        "batch_size": 1024,
    },
    # ============ Local SentenceTransformers ============
    # Large batches for GPU acceleration
    "minilm_384": {
        "model_id": "all-MiniLM-L6-v2",
        "dimension": 384,
        "column_suffix": "minilm_384",
        "type": "local",
        "batch_size": 512,
    },
    "mpnet_768": {
        "model_id": "all-mpnet-base-v2",
        "dimension": 768,
        "column_suffix": "mpnet_768",
        "type": "local",
        "batch_size": 256,
    },
    "bge_768": {
        "model_id": "BAAI/bge-base-en-v1.5",
        "dimension": 768,
        "column_suffix": "bge_768",
        "type": "local",
        "batch_size": 256,
    },
}

# Preset source groups
SOURCE_GROUPS = {
    # Fast: Just Gemini (free) + local MiniLM
    "minimal": ["gemini_256", "minilm_384"],
    # Standard: One per model family at 768-dim
    "standard": ["gemini_768", "openai_large_768", "mpnet_768"],
    # =========== ABLATION GROUPS (aligned dimensions) ===========
    # Model ablation at fixed dimension - compare models
    "ablation_256": ["gemini_256", "openai_large_256", "openai_small_256"],
    "ablation_512": ["gemini_512", "openai_large_512", "openai_small_512"],
    "ablation_768": ["gemini_768", "openai_large_768", "openai_small_768",
                     "mpnet_768", "bge_768"],
    # Dimension ablation per model - compare dimensions
    "ablation_gemini": ["gemini_256", "gemini_512", "gemini_768"],
    "ablation_openai_large": ["openai_large_256", "openai_large_512",
                              "openai_large_768", "openai_large_1536",
                              "openai_large_3072"],
    "ablation_openai_small": ["openai_small_256", "openai_small_512",
                              "openai_small_768", "openai_small_1536"],
    # Full ablation matrix (all aligned dims)
    "ablation_full": [
        # 256-dim row
        "gemini_256", "openai_large_256", "openai_small_256",
        # 512-dim row
        "gemini_512", "openai_large_512", "openai_small_512",
        # 768-dim row (includes local)
        "gemini_768", "openai_large_768", "openai_small_768",
        "mpnet_768", "bge_768",
        # Higher dims (OpenAI only)
        "openai_large_1536", "openai_large_3072", "openai_small_1536",
    ],
    # All: Everything including MiniLM
    "all": list(EMBEDDING_SOURCES.keys()),
    # Local-only: No API costs
    "local": ["minilm_384", "mpnet_768", "bge_768"],
    # Remote-only: Just API models
    "remote": [k for k, v in EMBEDDING_SOURCES.items() if v["type"] == "remote"],
}


def _text_hash(text: str) -> str:
    """Generate a unique hash for a text."""
    return hashlib.sha256(text.encode()).hexdigest()[:16]


def _estimate_tokens(text: str) -> int:
    """Rough token estimate: ~4 chars per token for English."""
    return len(text) // 4 + 1


def load_learned_batch_sizes(cache_dir: Path) -> Dict[str, int]:
    """Load previously learned optimal batch sizes."""
    path = cache_dir / "learned_batch_sizes.json"
    if path.exists():
        try:
            return json.loads(path.read_text())
        except Exception:
            return {}
    return {}


def save_learned_batch_size(cache_dir: Path, model_id: str, batch_size: int) -> None:
    """Persist learned optimal batch size for future runs."""
    path = cache_dir / "learned_batch_sizes.json"
    sizes = load_learned_batch_sizes(cache_dir)
    sizes[model_id] = batch_size
    with _cache_lock:
        path.write_text(json.dumps(sizes, indent=2))


def load_multi_embedding_cache(cache_dir: Path) -> Dict[str, Dict]:
    """Load multi-source embedding cache from disk.

    Structure: {source_name: {text_hash: embedding_list}}
    """
    cache_file = cache_dir / "multi_embedding_cache.json"
    if cache_file.exists():
        try:
            logger.info("Loading cache from %s", cache_file)
            with open(cache_file, "r") as f:
                cache = json.load(f)
            total_entries = sum(len(v) for v in cache.values())
            logger.info("Loaded cache with %d entries across %d sources",
                       total_entries, len(cache))
            return cache
        except json.JSONDecodeError:
            logger.warning("Cache file corrupted, starting fresh")
            return {}
    return {}


def save_multi_embedding_cache(cache: Dict[str, Dict], cache_dir: Path) -> None:
    """Save multi-source embedding cache to disk atomically."""
    cache_file = cache_dir / "multi_embedding_cache.json"
    temp_file = cache_file.with_suffix(".json.tmp")

    with _cache_lock:
        with open(temp_file, "w") as f:
            json.dump(cache, f)
        temp_file.rename(cache_file)


def _embed_single_batch(
    batch_texts: List[str],
    batch_indices: List[int],
    model_id: str,
    dimension: int,
    max_retries: int = 5,
) -> Tuple[List[Tuple[int, np.ndarray]], List[int], str]:
    """Embed a single batch with retries.

    Returns:
        Tuple of (results, failed_indices, error_type)
        error_type: "" (success), "rate_limit", "bad_request", "other"
    """
    from ember._internal.context.runtime import EmberContext
    from ember.models.runtime.registry import ModelRegistry

    registry = ModelRegistry(context=EmberContext.current())

    for attempt in range(max_retries):
        try:
            response = registry.embed_model(model_id, batch_texts, dimensions=dimension)
            results = [(idx, np.array(emb, dtype=np.float32))
                      for idx, emb in zip(batch_indices, response.embeddings)]
            return results, [], ""
        except Exception as e:
            error_str = str(e).lower()
            is_rate_limit = "rate" in error_str or "429" in error_str or "quota" in error_str
            is_bad_request = "400" in error_str or "bad request" in error_str or "too large" in error_str

            if is_bad_request:
                # Don't retry 400 errors - need smaller batch
                return [], batch_indices, "bad_request"
            elif is_rate_limit:
                backoff = min(60, 2 ** (attempt + 2))
                logger.warning("Rate limit hit, backing off %ds (attempt %d/%d)",
                              backoff, attempt + 1, max_retries)
                time.sleep(backoff)
            elif attempt == max_retries - 1:
                logger.error("Batch failed after %d attempts: %s", max_retries, e)
                return [], batch_indices, "other"
            else:
                time.sleep(2 ** attempt)
    return [], batch_indices, "rate_limit"


def compute_remote_embeddings_parallel(
    texts: List[str],
    model_id: str,
    dimension: int,
    batch_size: int = 256,
    max_workers: int = 64,
    max_retries: int = 5,
    cache_dir: Optional[Path] = None,
) -> Tuple[np.ndarray, List[int]]:
    """Compute embeddings with adaptive parallelism and binary-search batch sizing.

    Features:
    - Learns and persists optimal batch sizes across runs
    - Token-aware batching for outlier handling
    - Binary search for optimal batch size (converges fast)
    - Adaptive workers: reduces on rate limit errors
    - Canary batches for safe probing

    Args:
        texts: List of texts to embed.
        model_id: Ember model ID.
        dimension: Output dimension.
        batch_size: Initial texts per API call (adapts via binary search).
        max_workers: Initial concurrent requests (adapts down on rate limits).
        max_retries: Retry attempts per batch.
        cache_dir: Directory to persist learned batch sizes.

    Returns:
        Tuple of (embeddings array, list of failed indices)
    """
    all_embeddings = [None] * len(texts)
    final_failed = []
    current_workers = max_workers

    # Binary search state for batch size
    min_batch_size = 16
    max_batch_size = 2048

    # Check for previously learned optimal batch size
    learned_sizes = load_learned_batch_sizes(cache_dir) if cache_dir else {}
    if model_id in learned_sizes:
        learned = learned_sizes[model_id]
        current_batch_size = learned
        last_success_size = learned
        logger.info("Using learned batch size for %s: %d", model_id, learned)
    else:
        current_batch_size = batch_size
        last_success_size = min_batch_size

    last_fail_size = max_batch_size + 1  # Known to fail (or untested)
    stable_size = None  # Once we converge, stop adjusting

    # Track success for probing larger sizes
    consecutive_successes = 0
    probe_threshold = 20  # After N successes, try larger

    # Separate outliers (very long texts) to handle individually
    MAX_TOKENS_PER_TEXT = 8000  # OpenAI limit is 8191
    outlier_indices = []
    normal_texts = []

    for i, text in enumerate(texts):
        if _estimate_tokens(text) > MAX_TOKENS_PER_TEXT:
            outlier_indices.append(i)
            logger.warning("Text %d is very long (~%d tokens), will process individually",
                          i, _estimate_tokens(text))
        else:
            normal_texts.append((text, i))

    def make_batches_token_aware(
        texts_with_indices: List[Tuple[str, int]],
        max_size: int,
        max_tokens: int = 100000,
    ):
        """Create batches respecting both count and token limits."""
        batches = []
        current_batch = []
        current_tokens = 0

        for text, idx in texts_with_indices:
            text_tokens = _estimate_tokens(text)

            # Start new batch if limits exceeded
            if len(current_batch) >= max_size or current_tokens + text_tokens > max_tokens:
                if current_batch:
                    batches.append((
                        [t[0] for t in current_batch],
                        [t[1] for t in current_batch]
                    ))
                current_batch = []
                current_tokens = 0

            current_batch.append((text, idx))
            current_tokens += text_tokens

        # Don't forget last batch
        if current_batch:
            batches.append((
                [t[0] for t in current_batch],
                [t[1] for t in current_batch]
            ))

        return batches

    # Start with normal texts (outliers handled separately)
    remaining = normal_texts[:]
    round_num = 0
    max_rounds = 8

    while remaining and round_num < max_rounds:
        round_num += 1
        batches = make_batches_token_aware(remaining, current_batch_size)

        logger.info("  Round %d: %d texts in %d batches (size=%d, workers=%d)",
                   round_num, len(remaining), len(batches),
                   current_batch_size, current_workers)

        rate_limit_count = 0
        bad_request_count = 0
        success_count = 0
        next_remaining = []

        with ThreadPoolExecutor(max_workers=current_workers) as executor:
            futures = {
                executor.submit(
                    _embed_single_batch, batch_texts, batch_indices,
                    model_id, dimension, max_retries
                ): (batch_texts, batch_indices)
                for batch_texts, batch_indices in batches
            }

            for future in tqdm(as_completed(futures), total=len(futures),
                              desc=f"  {model_id}"):
                results, failed_indices, error_type = future.result()

                # Store successful embeddings
                for idx, emb in results:
                    all_embeddings[idx] = emb

                if error_type == "":
                    success_count += 1
                    consecutive_successes += 1
                elif error_type == "rate_limit":
                    rate_limit_count += 1
                    consecutive_successes = 0
                elif error_type == "bad_request":
                    bad_request_count += 1
                    consecutive_successes = 0

                # Queue failed indices for retry
                if failed_indices:
                    batch_texts, batch_indices = futures[future]
                    for idx in failed_indices:
                        pos = batch_indices.index(idx)
                        next_remaining.append((batch_texts[pos], idx))

        # Binary search batch size adjustment
        if stable_size is None:
            if bad_request_count > 0:
                # Current size too large - binary search down
                last_fail_size = current_batch_size
                old_size = current_batch_size
                current_batch_size = (last_success_size + current_batch_size) // 2
                current_batch_size = max(min_batch_size, current_batch_size)
                logger.warning("400 errors: binary search %d -> %d (range: %d-%d)",
                              old_size, current_batch_size, last_success_size, last_fail_size)

                # Check if converged
                if last_fail_size - last_success_size <= 32:
                    stable_size = last_success_size
                    current_batch_size = stable_size
                    logger.info("Batch size converged to %d", stable_size)

            elif success_count == len(batches):
                # All succeeded - this size works
                last_success_size = max(last_success_size, current_batch_size)

                # Probe larger with a single canary batch first
                if consecutive_successes >= probe_threshold and current_batch_size < last_fail_size - 32:
                    probe_size = (current_batch_size + last_fail_size) // 2
                    probe_size = min(max_batch_size, probe_size)

                    # Test canary batch with probe size
                    if remaining:
                        canary_batch = remaining[:probe_size]
                        canary_texts = [t[0] for t in canary_batch]
                        canary_indices = [t[1] for t in canary_batch]

                        logger.info("Probing with canary batch: size %d -> %d",
                                   current_batch_size, probe_size)
                        results, failed, error_type = _embed_single_batch(
                            canary_texts, canary_indices, model_id, dimension, 2
                        )

                        if error_type == "bad_request":
                            # Probe failed - update fail boundary, don't change size
                            last_fail_size = probe_size
                            logger.warning("Canary failed at %d, capping at %d",
                                          probe_size, current_batch_size)
                        elif error_type == "":
                            # Probe succeeded - scale up
                            for idx, emb in results:
                                all_embeddings[idx] = emb
                            # Remove canary batch from remaining
                            remaining = remaining[probe_size:]
                            current_batch_size = probe_size
                            last_success_size = probe_size
                            logger.info("Canary succeeded! Scaling to %d", probe_size)

                    consecutive_successes = 0

        # Adapt workers on rate limits
        if rate_limit_count > len(batches) * 0.1:
            old_workers = current_workers
            current_workers = max(4, current_workers // 2)
            logger.warning("Rate limits (%.0f%%), reducing workers %d -> %d",
                          rate_limit_count / len(batches) * 100,
                          old_workers, current_workers)
            time.sleep(5)

        remaining = next_remaining

    # Collect failures from normal texts
    final_failed = [idx for _, idx in remaining]

    # Handle outliers individually (batch size = 1)
    if outlier_indices:
        logger.info("Processing %d outlier texts individually", len(outlier_indices))
        for idx in tqdm(outlier_indices, desc="  Outliers"):
            text = texts[idx]
            results, failed, error_type = _embed_single_batch(
                [text], [idx], model_id, dimension, max_retries
            )
            if results:
                all_embeddings[idx] = results[0][1]
            else:
                final_failed.append(idx)
                logger.warning("Outlier text %d failed (too long: ~%d tokens)",
                              idx, _estimate_tokens(text))

    if final_failed:
        logger.warning("Failed to embed %d texts after %d rounds", len(final_failed), max_rounds)

    # Save learned batch size for future runs
    final_batch_size = stable_size or last_success_size
    if cache_dir and final_batch_size > min_batch_size:
        save_learned_batch_size(cache_dir, model_id, final_batch_size)
        logger.info("Saved learned batch size for %s: %d", model_id, final_batch_size)

    # Convert to array
    valid_embeddings = [e for e in all_embeddings if e is not None]
    if valid_embeddings:
        return np.array(valid_embeddings, dtype=np.float32), final_failed
    return np.array([]), final_failed


def compute_local_embeddings(
    texts: List[str],
    model_id: str,
    batch_size: int = 128,
) -> np.ndarray:
    """Compute embeddings using local SentenceTransformer."""
    from sentence_transformers import SentenceTransformer

    model = SentenceTransformer(model_id)
    embeddings = model.encode(
        texts,
        batch_size=batch_size,
        show_progress_bar=True,
        convert_to_numpy=True,
    )
    return embeddings


def compute_embeddings_for_source(
    texts: List[str],
    source_name: str,
    source_config: dict,
    cache: Dict[str, Dict],
    cache_dir: Path,
    max_workers: int = 8,
) -> Dict[str, List[float]]:
    """Compute embeddings for a single source with parallel processing.

    Args:
        texts: All texts to embed.
        source_name: Name of embedding source.
        source_config: Config dict for this source.
        cache: Multi-source cache dict (modified in place).
        cache_dir: Directory for cache persistence.
        max_workers: Concurrent API requests for remote sources.

    Returns:
        Dict mapping text_hash to embedding list.
    """
    source_cache = cache.get(source_name, {})

    # Find texts not in cache
    texts_to_compute = []
    hashes_to_compute = []
    for text in texts:
        h = _text_hash(text)
        if h not in source_cache:
            texts_to_compute.append(text)
            hashes_to_compute.append(h)

    if not texts_to_compute:
        logger.info("  %s: All %d texts cached", source_name, len(texts))
        return source_cache

    logger.info(
        "  %s: Computing %d/%d texts (%d cached) with %d workers",
        source_name,
        len(texts_to_compute),
        len(texts),
        len(texts) - len(texts_to_compute),
        max_workers if source_config["type"] == "remote" else 1,
    )

    batch_size = source_config.get("batch_size", 64)

    # Compute embeddings
    if source_config["type"] == "remote":
        # Parallel API calls for remote sources
        embeddings, failed = compute_remote_embeddings_parallel(
            texts_to_compute,
            source_config["model_id"],
            source_config["dimension"],
            batch_size=batch_size,
            max_workers=max_workers,
            cache_dir=cache_dir,
        )

        # Update cache with successful embeddings
        success_idx = 0
        for j, h in enumerate(hashes_to_compute):
            if j not in failed:
                source_cache[h] = embeddings[success_idx].tolist()
                success_idx += 1
    else:
        # Local models - compute all at once with internal batching
        embeddings = compute_local_embeddings(
            texts_to_compute,
            source_config["model_id"],
            batch_size,
        )
        for h, emb in zip(hashes_to_compute, embeddings):
            source_cache[h] = emb.tolist()

    # Final save
    cache[source_name] = source_cache
    save_multi_embedding_cache(cache, cache_dir)

    return source_cache


def compute_source_parallel(args: Tuple) -> Tuple[str, Dict]:
    """Worker function for parallel source computation."""
    source_name, texts, config, cache, cache_dir = args
    logger.info("Starting %s embeddings (%d-dim)...", source_name, config["dimension"])
    result = compute_embeddings_for_source(
        texts, source_name, config, cache, cache_dir
    )
    return source_name, result


def compute_all_embeddings(
    qa_path: Path,
    output_path: Path,
    cache_dir: Path,
    sources: List[str],
    parallel: int = 1,
    max_workers: int = 8,
    resume: bool = True,
) -> None:
    """Compute embeddings from multiple sources and save to parquet.

    Args:
        qa_path: Path to unified QA CSV.
        output_path: Path to save enriched parquet with embeddings.
        cache_dir: Directory for embedding cache.
        sources: List of source names to compute.
        parallel: Number of sources to compute in parallel (for local models).
        max_workers: Concurrent API requests per remote source.
        resume: Whether to resume from cache.
    """
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache = load_multi_embedding_cache(cache_dir) if resume else {}

    # Load data
    logger.info("Loading data from %s", qa_path)
    df = pd.read_csv(qa_path, low_memory=False)
    logger.info("Loaded %d rows", len(df))

    # Get unique texts
    all_queries = df["query"].drop_duplicates().tolist()
    all_task_descs = df["task_description"].drop_duplicates().tolist()
    all_texts = list(set(all_queries + all_task_descs))

    logger.info("Unique texts: %d queries, %d task_descriptions, %d total",
                len(all_queries), len(all_task_descs), len(all_texts))

    # Validate sources
    valid_sources = []
    for source in sources:
        if source in EMBEDDING_SOURCES:
            valid_sources.append(source)
        else:
            logger.warning("Unknown source: %s, skipping", source)

    # Separate local and remote sources (remote should be sequential to respect rate limits)
    local_sources = [s for s in valid_sources if EMBEDDING_SOURCES[s]["type"] == "local"]
    remote_sources = [s for s in valid_sources if EMBEDDING_SOURCES[s]["type"] == "remote"]

    # Process remote sources with parallel API calls
    for source_name in remote_sources:
        config = EMBEDDING_SOURCES[source_name]
        logger.info("Computing %s embeddings (%d-dim) with %d workers...",
                   source_name, config["dimension"], max_workers)
        compute_embeddings_for_source(all_texts, source_name, config, cache, cache_dir,
                                     max_workers=max_workers)

    # Process local sources in parallel if requested
    if local_sources and parallel > 1:
        logger.info("Computing %d local sources with %d parallel workers",
                   len(local_sources), parallel)

        with ThreadPoolExecutor(max_workers=parallel) as executor:
            futures = {}
            for source_name in local_sources:
                config = EMBEDDING_SOURCES[source_name]
                future = executor.submit(
                    compute_embeddings_for_source,
                    all_texts, source_name, config, cache, cache_dir
                )
                futures[future] = source_name

            for future in as_completed(futures):
                source_name = futures[future]
                try:
                    future.result()
                    logger.info("  %s completed", source_name)
                except Exception as e:
                    logger.error("  %s failed: %s", source_name, e)
    else:
        # Sequential processing
        for source_name in local_sources:
            config = EMBEDDING_SOURCES[source_name]
            logger.info("Computing %s embeddings (%d-dim)...", source_name, config["dimension"])
            compute_embeddings_for_source(all_texts, source_name, config, cache, cache_dir)

    # Build output dataframe with embedding columns
    logger.info("Building output dataframe...")

    for source_name in valid_sources:
        config = EMBEDDING_SOURCES[source_name]
        suffix = config["column_suffix"]
        source_cache = cache.get(source_name, {})

        query_col = f"query_embedding_{suffix}"
        task_col = f"task_description_embedding_{suffix}"

        df[query_col] = df["query"].apply(
            lambda q: str(source_cache.get(_text_hash(q), []))
        )
        df[task_col] = df["task_description"].apply(
            lambda t: str(source_cache.get(_text_hash(t), []))
        )

        logger.info("  Added columns: %s, %s", query_col, task_col)

    # Save output
    logger.info("Saving to %s", output_path)
    df.to_parquet(output_path, index=False)

    csv_path = output_path.with_suffix(".csv")
    df.to_csv(csv_path, index=False)
    logger.info("Saved CSV to %s", csv_path)

    # Print summary
    print("\n" + "=" * 70)
    print("MULTI-SOURCE EMBEDDING SUMMARY")
    print("=" * 70)
    print(f"\nRows: {len(df):,}")
    print(f"Sources computed: {valid_sources}")
    print("\nEmbedding columns:")
    for source in valid_sources:
        if source in EMBEDDING_SOURCES:
            suffix = EMBEDDING_SOURCES[source]["column_suffix"]
            dim = EMBEDDING_SOURCES[source]["dimension"]
            print(f"  - query_embedding_{suffix} ({dim}-dim)")
            print(f"  - task_description_embedding_{suffix} ({dim}-dim)")

    # Cache stats
    print("\nCache statistics:")
    for source_name, source_cache in cache.items():
        print(f"  - {source_name}: {len(source_cache):,} entries")


def main():
    parser = argparse.ArgumentParser(
        description="Compute embeddings from multiple sources with parallelism and resumption"
    )
    parser.add_argument(
        "--sources",
        type=str,
        default="minimal",
        help="Comma-separated sources or preset: minimal, standard, comprehensive, all",
    )
    parser.add_argument(
        "--parallel",
        type=int,
        default=4,
        help="Number of local models to compute in parallel (default: 4)",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=64,
        help="Concurrent API requests per remote source (default: 64)",
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        default=True,
        help="Resume from cache (default: True)",
    )
    parser.add_argument(
        "--no-resume",
        action="store_true",
        help="Start fresh, ignoring cache",
    )
    parser.add_argument(
        "--input",
        type=str,
        default=None,
        help="Input CSV path (default: from config)",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output parquet path (default: router_data_multi_embed.parquet)",
    )
    parser.add_argument(
        "--list-sources",
        action="store_true",
        help="List available embedding sources and exit",
    )
    args = parser.parse_args()

    if args.list_sources:
        print("\nAvailable Embedding Sources:")
        print("=" * 60)
        for name, config in EMBEDDING_SOURCES.items():
            print(f"  {name}:")
            print(f"    Model: {config['model_id']}")
            print(f"    Dimension: {config['dimension']}")
            print(f"    Type: {config['type']}")
            print(f"    Batch size: {config['batch_size']}")
            print()
        print("\nPreset Groups:")
        for group, sources in SOURCE_GROUPS.items():
            print(f"  {group}: {', '.join(sources)}")
        return

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

    # Resolve paths
    if args.input:
        qa_path = Path(args.input)
    else:
        qa_path = _MODULE_DIR / config.get("unified_qa_data_path", "data/unified_qa_data.csv")

    if args.output:
        output_path = Path(args.output)
    else:
        output_path = _MODULE_DIR / "data" / "router_data_multi_embed.parquet"

    cache_dir = _MODULE_DIR / "data" / ".router_cache"

    # Parse sources
    if args.sources in SOURCE_GROUPS:
        sources = SOURCE_GROUPS[args.sources]
    else:
        sources = [s.strip() for s in args.sources.split(",")]

    logger.info("Sources to compute: %s", sources)
    logger.info("Local parallel workers: %d", args.parallel)
    logger.info("Remote API concurrency: %d", args.workers)
    logger.info("Resume from cache: %s", not args.no_resume)

    compute_all_embeddings(
        qa_path=qa_path,
        output_path=output_path,
        cache_dir=cache_dir,
        sources=sources,
        parallel=args.parallel,
        max_workers=args.workers,
        resume=not args.no_resume,
    )


if __name__ == "__main__":
    import sys
    sys.path.insert(0, str(Path(__file__).resolve().parent))
    os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
    main()
