#!/usr/bin/env python3
"""
Generate benchmark data for dense and sparse vector indexes.

Downloads MS MARCO passages (or uses cached), generates:
- Dense embeddings via Triton server (jinaai/jina-embeddings-v3, Matryoshka-capable)
- Sparse vectors using naver/splade-v3 (local)

Usage:
    python generate_benchmark_data.py --num-docs 100000 --num-queries 1000

Output files:
    - dense_embeddings.bin: Dense vectors for documents (1024-dim, Jina v3)
    - dense_queries.bin: Dense vectors for queries
    - sparse_embeddings.bin: Sparse vectors for documents
    - sparse_queries.bin: Sparse vectors for queries
    - ground_truth_dense_full.bin: Ground truth for full-dim dense search
    - ground_truth_sparse.bin: Ground truth for sparse search

Requires:
    pip install tritonclient[grpc] transformers datasets numpy tqdm torch
"""

import argparse
import asyncio
import os
import struct
import sys
from pathlib import Path
from typing import Optional

import numpy as np
from tqdm import tqdm

try:
    import torch
    from transformers import AutoModelForMaskedLM, AutoTokenizer
except ImportError:
    print("Please install: pip install torch transformers")
    sys.exit(1)

try:
    from datasets import load_dataset
except ImportError:
    print("Please install: pip install datasets")
    sys.exit(1)


# ============================================================================
# Triton Client for Remote Embeddings
# ============================================================================


class TritonEmbedClient:
    """Async client for Triton embedding model (jina-v3-ensemble) via gRPC.

    Jina v3 outputs 1024-dim embeddings (Matryoshka-capable).
    Requires TRITON_URL and TRITON_API_KEY environment variables.
    """

    TASK_IDS = {
        "retrieval.query": 0,
        "retrieval.passage": 1,
        "separation": 2,
        "classification": 3,
        "text-matching": 4,
    }

    TASK_INSTRUCTIONS = {
        "retrieval.query": "Represent the query for retrieving evidence documents: ",
        "retrieval.passage": "Represent the document for retrieval: ",
    }

    def __init__(
        self,
        triton_url: str = None,
        api_key: str = None,
        model_name: str = "jina-v3-ensemble",
        tokenizer_name: str = "jinaai/jina-embeddings-v3",
    ):
        try:
            import tritonclient.grpc.aio as grpcclient

            self._grpcclient = grpcclient
        except ImportError:
            print("Please install: pip install tritonclient[grpc] grpcio")
            sys.exit(1)

        self.triton_url = triton_url or os.environ.get("TRITON_URL")
        if not self.triton_url:
            print("ERROR: TRITON_URL environment variable not set")
            sys.exit(1)

        self.api_key = api_key or os.environ.get("TRITON_API_KEY")
        if not self.api_key:
            print("ERROR: TRITON_API_KEY environment variable not set")
            sys.exit(1)

        self.model_name = model_name
        self._tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_name, trust_remote_code=True
        )
        self._headers = {"x-api-key": self.api_key}
        self._client = None
        print(f"Triton client initialized: {self.triton_url}")

    def _get_channel_args(self):
        """Get gRPC channel args with retry policy."""
        import json

        retry_config = json.dumps(
            {
                "methodConfig": [
                    {
                        "name": [{}],  # Apply to all methods
                        "retryPolicy": {
                            "maxAttempts": 5,
                            "initialBackoff": "0.5s",
                            "maxBackoff": "30s",
                            "backoffMultiplier": 2,
                            "retryableStatusCodes": [
                                "UNAVAILABLE",
                                "UNKNOWN",
                                "RESOURCE_EXHAUSTED",
                            ],
                        },
                    }
                ]
            }
        )
        return [("grpc.service_config", retry_config)]

    async def _get_client(self):
        """Get or create the gRPC client."""
        if self._client is None:
            self._client = self._grpcclient.InferenceServerClient(
                url=self.triton_url,
                ssl=False,
                channel_args=self._get_channel_args(),
            )
        return self._client

    async def encode(
        self,
        texts: list[str],
        task: str = "retrieval.passage",
        max_batch_size: int = 32,
    ) -> np.ndarray:
        """Encode texts to embeddings using Triton server."""
        client = await self._get_client()
        grpcclient = self._grpcclient

        all_embeddings = []
        for i in range(0, len(texts), max_batch_size):
            batch = texts[i : i + max_batch_size]

            # Prepend task instruction
            instruction = self.TASK_INSTRUCTIONS.get(task, "")
            prefixed_texts = [instruction + t for t in batch]

            encoded = self._tokenizer(
                prefixed_texts,
                padding=True,
                truncation=True,
                max_length=768,
                return_tensors="np",
            )

            input_ids = encoded["input_ids"].astype(np.int64)
            attention_mask = encoded["attention_mask"].astype(np.int64)
            task_id = np.array(
                [[self.TASK_IDS.get(task, 1)]] * len(batch), dtype=np.int64
            )

            inputs = [
                grpcclient.InferInput("input_ids", input_ids.shape, "INT64"),
                grpcclient.InferInput("attention_mask", attention_mask.shape, "INT64"),
                grpcclient.InferInput("task_id", task_id.shape, "INT64"),
            ]
            inputs[0].set_data_from_numpy(input_ids)
            inputs[1].set_data_from_numpy(attention_mask)
            inputs[2].set_data_from_numpy(task_id)

            outputs = [grpcclient.InferRequestedOutput("embedding")]
            result = await client.infer(
                model_name=self.model_name,
                inputs=inputs,
                outputs=outputs,
                headers=self._headers,
            )
            all_embeddings.append(result.as_numpy("embedding"))

        return np.vstack(all_embeddings)


async def generate_dense_embeddings_triton(
    texts: list[str],
    batch_size: int = 32,
    is_query: bool = False,
    output_path: Optional[Path] = None,
    checkpoint_interval: int = 1000,
) -> np.ndarray:
    """Generate dense embeddings using remote Triton server.

    Jina v3 on Triton outputs 1024-dim embeddings.
    Requires TRITON_URL and TRITON_API_KEY environment variables.
    Supports resuming from checkpoint if interrupted.
    """
    # Check for existing checkpoint
    start_idx = 0
    existing_embeddings = None
    if output_path:
        existing_embeddings, start_idx = load_dense_checkpoint(output_path)
        if start_idx > 0:
            print(
                f"Resuming from checkpoint: {start_idx}/{len(texts)} texts already processed"
            )

    if start_idx >= len(texts):
        print("All texts already processed, using cached embeddings")
        return existing_embeddings.astype(np.float32)

    client = TritonEmbedClient()
    task = "retrieval.query" if is_query else "retrieval.passage"

    remaining_texts = texts[start_idx:]
    print(
        f"Generating dense embeddings via Triton for {len(remaining_texts)} texts (task={task})..."
    )

    all_new_embeddings = []

    for i in tqdm(range(0, len(remaining_texts), batch_size), desc="Triton encoding"):
        batch = remaining_texts[i : i + batch_size]
        batch_embeddings = await client.encode(
            batch, task=task, max_batch_size=batch_size
        )
        all_new_embeddings.append(batch_embeddings)

        # Save checkpoint periodically
        processed = start_idx + i + len(batch)
        if output_path and (i + batch_size) % checkpoint_interval < batch_size:
            partial = np.vstack(all_new_embeddings)
            if existing_embeddings is not None:
                partial = np.vstack([existing_embeddings, partial])
            save_dense_checkpoint(partial, processed, output_path)

    # Combine all embeddings
    new_embeddings = np.vstack(all_new_embeddings)
    if existing_embeddings is not None:
        embeddings = np.vstack([existing_embeddings, new_embeddings])
    else:
        embeddings = new_embeddings

    # Normalize embeddings
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    embeddings = embeddings / np.maximum(norms, 1e-10)

    print(
        f"Generated {embeddings.shape[0]} vectors with {embeddings.shape[1]} dimensions"
    )
    return embeddings.astype(np.float32)


# ============================================================================
# Corpus Loading (MS MARCO Passage Ranking with qrels)
# ============================================================================


def load_msmarco_data(
    num_docs: int, num_queries: int, cache_dir: Optional[Path] = None
) -> tuple[list[str], list[str], list[str], list[str], dict[int, list[int]]]:
    """Load MS MARCO passage ranking dataset with qrels.

    Strategy:
    1. Load ALL qrels (train + dev) to get more queries
    2. Select exactly num_queries queries that have qrels
    3. Collect all their relevant passages
    4. Fill remaining with random passages to reach num_docs total

    Returns:
        passages: list of passage texts (num_docs total, includes relevant ones)
        passage_ids: list of passage IDs
        queries: list of query texts (exactly num_queries)
        query_ids: list of query IDs
        qrels: dict mapping query_idx -> list of relevant passage_idxs
    """
    print("Loading MS MARCO passage ranking dataset...")
    print(f"  Target: {num_queries} queries, {num_docs} passages")

    # Load ALL qrels (train has more data)
    print("Loading qrels (train + dev)...")
    qrels_dataset = load_dataset(
        "mteb/msmarco-v2",
        "default",
        cache_dir=str(cache_dir) if cache_dir else None,
        trust_remote_code=True,
    )

    # Build qrels dict from all splits
    print("Building qrels mapping...")
    qrels_raw: dict[str, list[str]] = {}
    for split_name in ["train", "dev", "dev2"]:
        if split_name in qrels_dataset:
            for item in tqdm(
                qrels_dataset[split_name], desc=f"Processing {split_name}"
            ):
                qid = item["query-id"]
                pid = item["corpus-id"]
                if qid not in qrels_raw:
                    qrels_raw[qid] = []
                qrels_raw[qid].append(pid)

    print(f"  Found {len(qrels_raw)} queries with relevance judgments")

    # Load queries
    print("Loading queries...")
    queries_dataset = load_dataset(
        "mteb/msmarco-v2",
        "queries",
        cache_dir=str(cache_dir) if cache_dir else None,
        trust_remote_code=True,
    )["queries"]

    # Build query lookup (only for queries with qrels)
    query_texts: dict[str, str] = {}
    for item in tqdm(queries_dataset, desc="Indexing queries"):
        qid = item["_id"]
        if qid in qrels_raw:
            query_texts[qid] = item["text"]

    # Select exactly num_queries queries
    selected_query_ids = [qid for qid in qrels_raw if qid in query_texts][:num_queries]
    print(f"Selected {len(selected_query_ids)} queries")

    # Collect all passage IDs needed for these queries
    needed_passage_ids: set[str] = set()
    for qid in selected_query_ids:
        for pid in qrels_raw[qid]:
            needed_passage_ids.add(pid)

    print(f"Need {len(needed_passage_ids)} relevant passages for selected queries")

    # Load corpus
    print("Loading corpus...")
    corpus = load_dataset(
        "mteb/msmarco-v2",
        "corpus",
        cache_dir=str(cache_dir) if cache_dir else None,
        trust_remote_code=True,
    )["corpus"]

    # Collect passages: first the needed ones, then random ones
    print("Extracting passages...")
    needed_passages: dict[str, str] = {}
    random_passages: list[tuple[str, str]] = []

    import random

    random.seed(42)

    # Calculate how many random passages we need
    num_random_needed = num_docs - len(needed_passage_ids)
    random_sample_rate = (num_random_needed * 2) / 138_000_000  # 2x oversample

    for item in tqdm(corpus, desc="Loading passages"):
        pid = item["_id"]
        text = item.get("text", "") or item.get("title", "")
        if not text:
            continue

        if pid in needed_passage_ids:
            needed_passages[pid] = text
        elif random.random() < random_sample_rate:
            random_passages.append((pid, text))

        # Stop early if we have everything
        if (
            len(needed_passages) >= len(needed_passage_ids)
            and len(random_passages) >= num_random_needed
        ):
            break

    print(f"  Collected {len(needed_passages)} relevant passages")
    print(f"  Collected {len(random_passages)} random passages")

    # Build final passage list: relevant passages first, then random
    passages = []
    passage_ids = []
    passage_id_to_idx: dict[str, int] = {}

    # Add relevant passages first
    for pid, text in needed_passages.items():
        passage_id_to_idx[pid] = len(passages)
        passage_ids.append(pid)
        passages.append(text)

    # Fill up to num_docs with random passages
    for pid, text in random_passages:
        if len(passages) >= num_docs:
            break
        if pid not in passage_id_to_idx:  # Avoid duplicates
            passage_id_to_idx[pid] = len(passages)
            passage_ids.append(pid)
            passages.append(text)

    print(
        f"Total passages: {len(passages)} ({len(needed_passages)} relevant + {len(passages) - len(needed_passages)} random)"
    )

    # Build final queries and qrels
    queries = []
    query_ids_out = []
    valid_qrels: dict[int, list[int]] = {}

    for qid in selected_query_ids:
        if qid not in query_texts:
            continue

        relevant_idxs = []
        for pid in qrels_raw[qid]:
            if pid in passage_id_to_idx:
                relevant_idxs.append(passage_id_to_idx[pid])

        if relevant_idxs:
            query_idx = len(queries)
            queries.append(query_texts[qid])
            query_ids_out.append(qid)
            valid_qrels[query_idx] = relevant_idxs

    print(f"Final: {len(queries)} queries with valid qrels")
    avg_relevant = (
        sum(len(v) for v in valid_qrels.values()) / len(valid_qrels)
        if valid_qrels
        else 0
    )
    print(f"  Average relevant passages per query: {avg_relevant:.1f}")

    return passages, passage_ids, queries, query_ids_out, valid_qrels


def load_msmarco_passages(num_docs: int, cache_dir: Optional[Path] = None) -> list[str]:
    """Load MS MARCO passages dataset (legacy interface)."""
    print(f"Loading MS MARCO passages (requesting {num_docs} documents)...")

    # Load the passage corpus
    dataset = load_dataset(
        "microsoft/ms_marco",
        "v1.1",
        split="train",
        cache_dir=str(cache_dir) if cache_dir else None,
        trust_remote_code=True,
    )

    # Extract passages - MS MARCO has nested structure
    passages = []
    seen = set()

    for item in tqdm(dataset, desc="Extracting passages"):
        if len(passages) >= num_docs:
            break
        # Each item has 'passages' which is a dict with 'passage_text' list
        if "passages" in item and "passage_text" in item["passages"]:
            for text in item["passages"]["passage_text"]:
                if text and text not in seen:
                    seen.add(text)
                    passages.append(text)
                    if len(passages) >= num_docs:
                        break

    print(f"Loaded {len(passages)} unique passages")
    return passages


def load_msmarco_queries(
    num_queries: int, cache_dir: Optional[Path] = None
) -> list[str]:
    """Load MS MARCO queries (legacy interface)."""
    print(f"Loading MS MARCO queries (requesting {num_queries})...")

    dataset = load_dataset(
        "microsoft/ms_marco",
        "v1.1",
        split="train",
        cache_dir=str(cache_dir) if cache_dir else None,
        trust_remote_code=True,
    )

    queries = []
    seen = set()

    for item in tqdm(dataset, desc="Extracting queries"):
        if len(queries) >= num_queries:
            break
        query = item.get("query", "")
        if query and query not in seen:
            seen.add(query)
            queries.append(query)

    print(f"Loaded {len(queries)} unique queries")
    return queries


def truncate_embeddings(
    embeddings: np.ndarray, target_dim: int, silent: bool = False
) -> np.ndarray:
    """Apply Matryoshka truncation to embeddings.

    Jina v3 supports dimensions: 32, 64, 128, 256, 512, 768, 1024
    """
    if embeddings.shape[1] <= target_dim:
        return embeddings

    if not silent:
        print(
            f"Truncating from {embeddings.shape[1]} to {target_dim} dimensions (Matryoshka)"
        )
    truncated = embeddings[:, :target_dim]
    # Re-normalize after truncation
    norms = np.linalg.norm(truncated, axis=1, keepdims=True)
    truncated = truncated / np.maximum(norms, 1e-10)
    return truncated.astype(np.float32)


# ============================================================================
# Sparse Embeddings (SPLADE-v3)
# ============================================================================


def get_device() -> str:
    """Get the best available device for PyTorch."""
    if torch.cuda.is_available():
        return "cuda"
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


class SpladeEncoder:
    """SPLADE-v3 sparse encoder with GPU/MPS optimization."""

    def __init__(self, model_name: str = "naver/splade-v3"):
        print(f"Loading SPLADE model: {model_name}")
        self.device = get_device()

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForMaskedLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16
            if self.device in ("cuda", "mps")
            else torch.float32,
        )
        self.model.to(self.device)
        self.model.eval()

        print(f"SPLADE model loaded on {self.device} (fp16: {self.device != 'cpu'})")

    def encode(
        self, texts: list[str], batch_size: int = 32
    ) -> list[tuple[np.ndarray, np.ndarray]]:
        """Encode texts to sparse vectors.

        Returns list of (indices, values) tuples.
        """
        results = []

        for i in tqdm(range(0, len(texts), batch_size), desc="SPLADE encoding"):
            batch = texts[i : i + batch_size]
            batch_results = self._encode_batch(batch)
            results.extend(batch_results)

        return results

    def _encode_batch(self, texts: list[str]) -> list[tuple[np.ndarray, np.ndarray]]:
        """Encode a batch of texts."""
        inputs = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt",
        ).to(self.device)

        with torch.no_grad():
            outputs = self.model(**inputs)
            # SPLADE uses log(1 + ReLU(x)) aggregated over tokens
            logits = outputs.logits
            # Max pooling over sequence length with attention mask
            attention_mask = inputs["attention_mask"].unsqueeze(-1)
            logits = logits * attention_mask
            # Apply SPLADE activation: log(1 + ReLU(x))
            splade_rep = torch.log1p(torch.relu(logits))
            # Max pool over sequence
            splade_rep = torch.max(splade_rep, dim=1).values

        results = []
        for rep in splade_rep:
            # Get non-zero indices and values
            rep_cpu = rep.cpu().numpy()
            nonzero_mask = rep_cpu > 0.0
            indices = np.where(nonzero_mask)[0].astype(np.uint32)
            values = rep_cpu[nonzero_mask].astype(np.float32)

            # Sort by index for consistent ordering
            sort_idx = np.argsort(indices)
            indices = indices[sort_idx]
            values = values[sort_idx]

            results.append((indices, values))

        return results


def generate_sparse_embeddings(
    texts: list[str],
    model_name: str = "naver/splade-v3",
    batch_size: int = 32,
    output_path: Optional[Path] = None,
    checkpoint_interval: int = 1000,
) -> list[tuple[np.ndarray, np.ndarray]]:
    """Generate sparse embeddings using SPLADE-v3.

    Supports resuming from checkpoint if interrupted.
    """
    # Check for existing checkpoint
    start_idx = 0
    existing_vectors: list[tuple[np.ndarray, np.ndarray]] = []
    if output_path:
        loaded, start_idx = load_sparse_checkpoint(output_path)
        if loaded:
            existing_vectors = loaded
            print(
                f"Resuming from checkpoint: {start_idx}/{len(texts)} texts already processed"
            )

    if start_idx >= len(texts):
        print("All texts already processed, using cached embeddings")
        return existing_vectors

    encoder = SpladeEncoder(model_name)
    remaining_texts = texts[start_idx:]

    results = list(existing_vectors)

    for i in tqdm(range(0, len(remaining_texts), batch_size), desc="SPLADE encoding"):
        batch = remaining_texts[i : i + batch_size]
        batch_results = encoder._encode_batch(batch)
        results.extend(batch_results)

        # Save checkpoint periodically
        processed = start_idx + i + len(batch)
        if output_path and (i + batch_size) % checkpoint_interval < batch_size:
            save_sparse_checkpoint(results, processed, output_path)

    return results


# ============================================================================
# Ground Truth Computation
# ============================================================================


def compute_dense_ground_truth(
    query_embeddings: np.ndarray,
    doc_embeddings: np.ndarray,
    k: int = 100,
) -> np.ndarray:
    """Compute ground truth for dense vectors using brute force."""
    print(f"Computing dense ground truth (k={k})...")

    num_queries = query_embeddings.shape[0]
    ground_truth = np.zeros((num_queries, k), dtype=np.uint32)

    for i in tqdm(range(num_queries), desc="Dense ground truth"):
        # Compute dot product (vectors are normalized, so this is cosine similarity)
        scores = np.dot(doc_embeddings, query_embeddings[i])
        # Get top-k indices
        top_k = np.argsort(-scores)[:k]
        ground_truth[i] = top_k

    return ground_truth


def compute_sparse_ground_truth(
    query_vectors: list[tuple[np.ndarray, np.ndarray]],
    doc_vectors: list[tuple[np.ndarray, np.ndarray]],
    k: int = 100,
) -> np.ndarray:
    """Compute ground truth for sparse vectors using brute force dot product."""
    print(f"Computing sparse ground truth (k={k})...")

    num_queries = len(query_vectors)
    num_docs = len(doc_vectors)
    ground_truth = np.zeros((num_queries, k), dtype=np.uint32)

    for qi in tqdm(range(num_queries), desc="Sparse ground truth"):
        q_indices, q_values = query_vectors[qi]
        scores = np.zeros(num_docs, dtype=np.float32)

        for di in range(num_docs):
            d_indices, d_values = doc_vectors[di]
            # Compute sparse dot product
            score = sparse_dot(q_indices, q_values, d_indices, d_values)
            scores[di] = score

        # Get top-k indices
        top_k = np.argsort(-scores)[:k]
        ground_truth[qi] = top_k

    return ground_truth


def sparse_dot(
    indices1: np.ndarray,
    values1: np.ndarray,
    indices2: np.ndarray,
    values2: np.ndarray,
) -> float:
    """Compute dot product of two sparse vectors."""
    i, j = 0, 0
    result = 0.0

    while i < len(indices1) and j < len(indices2):
        if indices1[i] == indices2[j]:
            result += values1[i] * values2[j]
            i += 1
            j += 1
        elif indices1[i] < indices2[j]:
            i += 1
        else:
            j += 1

    return result


# ============================================================================
# File I/O with Checkpoint Support
# ============================================================================


def get_checkpoint_path(path: Path) -> Path:
    """Get checkpoint file path for a given output file."""
    return path.with_suffix(path.suffix + ".checkpoint.npz")


def save_dense_checkpoint(embeddings: np.ndarray, processed: int, path: Path):
    """Save dense embeddings checkpoint."""
    checkpoint_path = get_checkpoint_path(path)
    np.savez(checkpoint_path, embeddings=embeddings, processed=processed)


def load_dense_checkpoint(path: Path) -> tuple[np.ndarray | None, int]:
    """Load dense embeddings checkpoint if exists."""
    checkpoint_path = get_checkpoint_path(path)
    if checkpoint_path.exists():
        data = np.load(checkpoint_path)
        return data["embeddings"], int(data["processed"])
    return None, 0


def save_sparse_checkpoint(
    vectors: list[tuple[np.ndarray, np.ndarray]], processed: int, path: Path
):
    """Save sparse embeddings checkpoint."""
    checkpoint_path = get_checkpoint_path(path)
    # Convert to arrays for saving
    indices_list = [v[0] for v in vectors]
    values_list = [v[1] for v in vectors]
    np.savez(
        checkpoint_path,
        indices=np.array(indices_list, dtype=object),
        values=np.array(values_list, dtype=object),
        processed=processed,
    )


def load_sparse_checkpoint(
    path: Path,
) -> tuple[list[tuple[np.ndarray, np.ndarray]] | None, int]:
    """Load sparse embeddings checkpoint if exists."""
    checkpoint_path = get_checkpoint_path(path)
    if checkpoint_path.exists():
        data = np.load(checkpoint_path, allow_pickle=True)
        indices_list = data["indices"]
        values_list = data["values"]
        vectors = [(idx, val) for idx, val in zip(indices_list, values_list)]
        return vectors, int(data["processed"])
    return None, 0


def cleanup_checkpoint(path: Path):
    """Remove checkpoint file after successful completion."""
    checkpoint_path = get_checkpoint_path(path)
    if checkpoint_path.exists():
        checkpoint_path.unlink()
        print(f"  Cleaned up checkpoint: {checkpoint_path}")


def save_dense_embeddings(embeddings: np.ndarray, path: Path):
    """Save dense embeddings in binary format.

    Format: num_vectors (u32), dim (u32), vectors (f32 * num_vectors * dim)
    """
    num_vectors, dim = embeddings.shape

    with open(path, "wb") as f:
        f.write(struct.pack("<II", num_vectors, dim))
        f.write(embeddings.tobytes())

    print(f"Saved {num_vectors} dense vectors (dim={dim}) to {path}")
    print(f"  File size: {path.stat().st_size / 1024 / 1024:.2f} MB")
    cleanup_checkpoint(path)


def load_dense_embeddings(path: Path) -> np.ndarray:
    """Load dense embeddings from binary format."""
    with open(path, "rb") as f:
        num_vectors, dim = struct.unpack("<II", f.read(8))
        data = np.frombuffer(f.read(), dtype=np.float32)
        embeddings = data.reshape(num_vectors, dim)
    print(f"  Loaded {num_vectors} vectors (dim={dim})")
    return embeddings


def save_sparse_embeddings(vectors: list[tuple[np.ndarray, np.ndarray]], path: Path):
    """Save sparse embeddings in binary format.

    Format:
    - num_vectors: u32
    - For each vector:
        - num_nonzero: u32
        - indices: [u32; num_nonzero]
        - values: [f32; num_nonzero]
    """
    with open(path, "wb") as f:
        f.write(struct.pack("<I", len(vectors)))

        total_nnz = 0
        for indices, values in vectors:
            num_nonzero = len(indices)
            total_nnz += num_nonzero
            f.write(struct.pack("<I", num_nonzero))
            f.write(indices.astype(np.uint32).tobytes())
            f.write(values.astype(np.float32).tobytes())

    avg_nnz = total_nnz / len(vectors) if vectors else 0
    print(f"Saved {len(vectors)} sparse vectors to {path}")
    print(f"  Avg non-zeros: {avg_nnz:.1f}")
    print(f"  File size: {path.stat().st_size / 1024 / 1024:.2f} MB")
    cleanup_checkpoint(path)


def load_sparse_embeddings(path: Path) -> list[tuple[np.ndarray, np.ndarray]]:
    """Load sparse embeddings from binary format."""
    vectors = []
    with open(path, "rb") as f:
        num_vectors = struct.unpack("<I", f.read(4))[0]
        for _ in range(num_vectors):
            num_nonzero = struct.unpack("<I", f.read(4))[0]
            indices = np.frombuffer(f.read(num_nonzero * 4), dtype=np.uint32).copy()
            values = np.frombuffer(f.read(num_nonzero * 4), dtype=np.float32).copy()
            vectors.append((indices, values))
    print(f"  Loaded {num_vectors} sparse vectors")
    return vectors


def save_ground_truth(ground_truth: np.ndarray, path: Path):
    """Save ground truth in binary format.

    Format: num_queries (u32), k (u32), indices (u32 * num_queries * k)
    """
    num_queries, k = ground_truth.shape

    with open(path, "wb") as f:
        f.write(struct.pack("<II", num_queries, k))
        f.write(ground_truth.astype(np.uint32).tobytes())

    print(f"Saved ground truth ({num_queries} queries, k={k}) to {path}")


def save_qrels(qrels: dict[int, list[int]], path: Path):
    """Save qrels (relevance judgments) in binary format.

    Format:
    - num_queries: u32
    - For each query:
        - query_idx: u32
        - num_relevant: u32
        - relevant_passage_idxs: [u32; num_relevant]
    """
    with open(path, "wb") as f:
        f.write(struct.pack("<I", len(qrels)))
        for query_idx, relevant_idxs in qrels.items():
            f.write(struct.pack("<II", query_idx, len(relevant_idxs)))
            f.write(np.array(relevant_idxs, dtype=np.uint32).tobytes())

    total_relevant = sum(len(v) for v in qrels.values())
    print(
        f"Saved qrels ({len(qrels)} queries, {total_relevant} total relevant) to {path}"
    )


def load_qrels(path: Path) -> dict[int, list[int]]:
    """Load qrels from binary format."""
    qrels = {}
    with open(path, "rb") as f:
        num_queries = struct.unpack("<I", f.read(4))[0]
        for _ in range(num_queries):
            query_idx, num_relevant = struct.unpack("<II", f.read(8))
            relevant_idxs = list(
                np.frombuffer(f.read(num_relevant * 4), dtype=np.uint32)
            )
            qrels[query_idx] = relevant_idxs
    print(f"  Loaded qrels for {len(qrels)} queries")
    return qrels


# ============================================================================
# Main
# ============================================================================


async def async_main(args):
    """Async main function."""
    # Pre-download tokenizers/models
    print("\n" + "=" * 60)
    print("Initializing")
    print("=" * 60)

    if not args.skip_dense:
        triton_url = os.environ.get("TRITON_URL")
        if not triton_url:
            print("ERROR: Set TRITON_URL environment variable")
            sys.exit(1)
        if not os.environ.get("TRITON_API_KEY"):
            print("ERROR: Set TRITON_API_KEY environment variable")
            sys.exit(1)
        print(f"Dense embeddings via Triton: {triton_url}")
        AutoTokenizer.from_pretrained(
            "jinaai/jina-embeddings-v3", trust_remote_code=True
        )
        print("  Tokenizer ready")

    if not args.skip_sparse:
        print(f"Downloading sparse model: {args.sparse_model}")
        AutoTokenizer.from_pretrained(args.sparse_model)
        AutoModelForMaskedLM.from_pretrained(args.sparse_model)
        print("  Sparse model ready")

    print("Ready!\n")

    # Load corpus
    qrels = None
    if args.use_beir:
        passages, _, queries, _, qrels = load_msmarco_data(
            args.num_docs, args.num_queries, args.cache_dir
        )
        # Save qrels for benchmark use
        qrels_path = args.output_dir / "qrels.bin"
        save_qrels(qrels, qrels_path)
    else:
        passages = load_msmarco_passages(args.num_docs, args.cache_dir)
        queries = load_msmarco_queries(args.num_queries, args.cache_dir)

    # Save raw texts for BM25 benchmarking
    texts_path = args.output_dir / "passages.txt"
    if not texts_path.exists():
        print(f"Saving {len(passages)} passages to {texts_path}")
        with open(texts_path, "w", encoding="utf-8") as f:
            for p in passages:
                # Replace newlines with spaces to keep one passage per line
                f.write(p.replace("\n", " ").replace("\r", "") + "\n")

    queries_text_path = args.output_dir / "queries.txt"
    if not queries_text_path.exists():
        print(f"Saving {len(queries)} queries to {queries_text_path}")
        with open(queries_text_path, "w", encoding="utf-8") as f:
            for q in queries:
                f.write(q.replace("\n", " ").replace("\r", "") + "\n")

    # Ensure we have enough data
    if len(passages) < args.num_docs:
        print(f"Warning: Only got {len(passages)} passages (requested {args.num_docs})")
    if len(queries) < args.num_queries:
        print(
            f"Warning: Only got {len(queries)} queries (requested {args.num_queries})"
        )

    # Generate dense embeddings
    if not args.skip_dense:
        print("\n" + "=" * 60)
        print("Generating Dense Embeddings")
        print("=" * 60)

        doc_output = args.output_dir / "dense_embeddings.bin"
        query_output = args.output_dir / "dense_queries.bin"

        # Check if embeddings already exist
        if doc_output.exists() and query_output.exists():
            print(f"Loading existing dense embeddings from {doc_output}")
            doc_embeddings = load_dense_embeddings(doc_output)
            print(f"Loading existing dense queries from {query_output}")
            query_embeddings = load_dense_embeddings(query_output)
        else:
            # Use Triton server for embeddings
            doc_embeddings = await generate_dense_embeddings_triton(
                passages,
                batch_size=args.batch_size,
                is_query=False,
                output_path=doc_output,
            )
            save_dense_embeddings(doc_embeddings, doc_output)

            query_embeddings = await generate_dense_embeddings_triton(
                queries,
                batch_size=args.batch_size,
                is_query=True,
                output_path=query_output,
            )
            save_dense_embeddings(query_embeddings, query_output)

        # Compute ground truth for full-dimensional embeddings (skip if exists)
        gt_full_path = args.output_dir / "ground_truth_dense_full.bin"
        if gt_full_path.exists():
            print(f"Full-dim ground truth already exists: {gt_full_path}")
        else:
            print(
                f"Computing ground truth for full {doc_embeddings.shape[1]}-dim embeddings..."
            )
            ground_truth_full = compute_dense_ground_truth(
                query_embeddings, doc_embeddings, args.k
            )
            save_ground_truth(ground_truth_full, gt_full_path)

        # Compute ground truth for truncated embeddings (skip if exists)
        gt_truncated_path = args.output_dir / f"ground_truth_dense_{args.dense_dim}.bin"
        if gt_truncated_path.exists():
            print(f"Truncated ground truth already exists: {gt_truncated_path}")
        else:
            print(
                f"Computing ground truth for truncated {args.dense_dim}-dim embeddings..."
            )
            doc_truncated = truncate_embeddings(
                doc_embeddings, args.dense_dim, silent=True
            )
            query_truncated = truncate_embeddings(
                query_embeddings, args.dense_dim, silent=True
            )
            ground_truth_truncated = compute_dense_ground_truth(
                query_truncated, doc_truncated, args.k
            )
            save_ground_truth(ground_truth_truncated, gt_truncated_path)

    # Generate sparse embeddings
    if not args.skip_sparse:
        print("\n" + "=" * 60)
        print("Generating Sparse Embeddings (SPLADE)")
        print("=" * 60)

        doc_sparse_output = args.output_dir / "sparse_embeddings.bin"
        query_sparse_output = args.output_dir / "sparse_queries.bin"

        # Check if sparse embeddings already exist
        if doc_sparse_output.exists() and query_sparse_output.exists():
            print(f"Loading existing sparse embeddings from {doc_sparse_output}")
            doc_sparse = load_sparse_embeddings(doc_sparse_output)
            print(f"Loading existing sparse queries from {query_sparse_output}")
            query_sparse = load_sparse_embeddings(query_sparse_output)
        else:
            doc_sparse = generate_sparse_embeddings(
                passages,
                model_name=args.sparse_model,
                batch_size=args.batch_size,
                output_path=doc_sparse_output,
            )
            save_sparse_embeddings(doc_sparse, doc_sparse_output)

            query_sparse = generate_sparse_embeddings(
                queries,
                model_name=args.sparse_model,
                batch_size=args.batch_size,
                output_path=query_sparse_output,
            )
            save_sparse_embeddings(query_sparse, query_sparse_output)

        # Compute ground truth (skip if already exists)
        sparse_gt_path = args.output_dir / "ground_truth_sparse.bin"
        if sparse_gt_path.exists():
            print(f"Sparse ground truth already exists: {sparse_gt_path}")
        else:
            ground_truth = compute_sparse_ground_truth(query_sparse, doc_sparse, args.k)
            save_ground_truth(ground_truth, sparse_gt_path)

    print("\n" + "=" * 60)
    print("Benchmark data generation complete!")
    print("=" * 60)
    print(f"\nOutput directory: {args.output_dir}")
    print("\nTo run benchmarks:")
    print(f"  BENCHMARK_DATA={args.output_dir} cargo bench --bench large_benchmark")


def main():
    parser = argparse.ArgumentParser(
        description="Generate benchmark data for vector indexes"
    )
    parser.add_argument(
        "--output-dir",
        "-o",
        type=Path,
        default=Path("benchmark_data"),
        help="Output directory",
    )
    parser.add_argument(
        "--num-docs",
        "-n",
        type=int,
        default=100000,
        help="Number of documents to index",
    )
    parser.add_argument(
        "--num-queries", "-q", type=int, default=10000, help="Number of queries"
    )
    parser.add_argument(
        "--dense-dim",
        "-d",
        type=int,
        default=128,
        help="Dense embedding dimension for ground truth (Matryoshka truncation)",
    )
    parser.add_argument(
        "--k", type=int, default=100, help="Number of ground truth neighbors"
    )
    parser.add_argument(
        "--sparse-model",
        type=str,
        default="naver/splade-v3",
        help="Sparse embedding model (local)",
    )
    parser.add_argument(
        "--batch-size",
        "-b",
        type=int,
        default=32,
        help="Batch size for Triton encoding",
    )
    parser.add_argument(
        "--skip-dense", action="store_true", help="Skip dense embedding generation"
    )
    parser.add_argument(
        "--skip-sparse", action="store_true", help="Skip sparse embedding generation"
    )
    parser.add_argument(
        "--cache-dir", type=Path, default=None, help="Cache directory for datasets"
    )
    parser.add_argument(
        "--use-beir",
        action="store_true",
        help="Use BeIR/msmarco dataset with official qrels for IR evaluation",
    )

    args = parser.parse_args()

    # Create output directory
    args.output_dir.mkdir(parents=True, exist_ok=True)

    asyncio.run(async_main(args))


if __name__ == "__main__":
    main()
