#!/usr/bin/env python3
"""
Full Index Test for CyborgDB Python SDK

This script tests the service's ability to handle large-scale vector indexing
using the wiki_all_1M dataset from ann-benchmarks:
1. Downloads the wiki_all_1M dataset (if not cached)
2. Creates an index
3. Inserts vectors in batches
4. Trains the index
5. Runs queries and measures performance
6. Measures recall against precomputed ground truth
7. Reports throughput and quality metrics
"""

import os
import sys
import time
import requests
import numpy as np
import h5py
from tqdm import tqdm

# Import from cyborgdb SDK
from cyborgdb import Client, IndexIVFFlat


DATASET_URL = "https://wiki-all.s3.us-east-1.amazonaws.com/wiki_all_1M.hdf5"
DATASET_FILENAME = "wiki_all_1M.hdf5"


def download_dataset(cache_dir: str) -> str:
    """Download the wiki_all_1M dataset if not already cached."""
    os.makedirs(cache_dir, exist_ok=True)
    dataset_path = os.path.join(cache_dir, DATASET_FILENAME)

    if os.path.exists(dataset_path):
        print(f"  Dataset already cached at {dataset_path}")
        return dataset_path

    print(f"  Downloading from {DATASET_URL}...")
    response = requests.get(DATASET_URL, stream=True)
    response.raise_for_status()

    total_size = int(response.headers.get("content-length", 0))

    with open(dataset_path, "wb") as f:
        with tqdm(
            total=total_size, unit="B", unit_scale=True, desc="Downloading"
        ) as pbar:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
                pbar.update(len(chunk))

    print(f"  Dataset downloaded to {dataset_path}")
    return dataset_path


def load_dataset(
    dataset_path: str,
    num_vectors: int = None,
    num_queries: int = 1000,
    top_k: int = 100,
):
    """Load the wiki_all_1M dataset from HDF5 file."""
    with h5py.File(dataset_path, "r") as f:
        print(f"  Available keys: {list(f.keys())}")

        # Load train vectors (the vectors to index)
        train_data = np.array(f["train"], dtype=np.float32)
        if num_vectors and num_vectors < len(train_data):
            train_data = train_data[:num_vectors]

        # Load test vectors (query vectors)
        test_data = np.array(f["test"], dtype=np.float32)
        if num_queries < len(test_data):
            test_data = test_data[:num_queries]

        # Load precomputed ground truth neighbors
        neighbors = np.array(f["neighbors"], dtype=np.int32)
        if num_queries < len(neighbors):
            neighbors = neighbors[:num_queries]
        if top_k < neighbors.shape[1]:
            neighbors = neighbors[:, :top_k]

    return train_data, test_data, neighbors


def compute_recall(ground_truth: np.ndarray, result_ids: list) -> float:
    """Compute recall: what fraction of ground truth results were found."""
    if len(ground_truth) == 0:
        return 0.0
    gt_set = set(ground_truth.tolist())
    result_set = set(int(rid) for rid in result_ids)
    return len(gt_set & result_set) / len(gt_set)


def main():
    # Configuration from environment
    base_url = os.environ.get("CYBORGDB_URL", "http://localhost:8000")
    api_key = os.environ.get("CYBORGDB_API_KEY")
    vector_count = int(os.environ.get("VECTOR_COUNT", "1000000"))  # Default to full 1M
    batch_size = int(os.environ.get("BATCH_SIZE", "10000"))
    top_k = int(os.environ.get("TOP_K", "100"))
    num_queries = int(os.environ.get("NUM_QUERIES", "1000"))
    recall_baseline = float(
        os.environ.get("RECALL_BASELINE", "0.95")
    )  # 95% for IVFFLAT index
    cache_dir = os.environ.get("DATASET_CACHE_DIR", "/tmp/ann-benchmarks")

    if not api_key:
        print("ERROR: CYBORGDB_API_KEY environment variable not set")
        sys.exit(1)

    print("=" * 60)
    print("CyborgDB Full Index Test - Python SDK (wiki_all_1M)")
    print("=" * 60)
    print(f"Service URL: {base_url}")
    print(f"Vector count: {vector_count:,}")
    print(f"Batch size: {batch_size:,}")
    print(f"Top K: {top_k}")
    print(f"Number of queries: {num_queries}")
    print(f"Recall baseline: {recall_baseline:.0%}")
    print("=" * 60)

    # Download dataset
    print("\n[1/8] Downloading/loading wiki_all_1M dataset...")
    dataset_path = download_dataset(cache_dir)

    # Load dataset
    print("\n[2/8] Loading dataset...")
    train_vectors, query_vectors, ground_truth = load_dataset(
        dataset_path, num_vectors=vector_count, num_queries=num_queries, top_k=top_k
    )
    vector_dim = train_vectors.shape[1]
    actual_vector_count = len(train_vectors)

    print(f"  Train vectors: {train_vectors.shape}")
    print(f"  Query vectors: {query_vectors.shape}")
    print(f"  Ground truth: {ground_truth.shape}")
    print(f"  Vector dimension: {vector_dim}")

    # Initialize client
    print("\n[3/8] Initializing client...")
    client = Client(base_url=base_url, api_key=api_key)

    # Check health
    health = client.get_health()
    print(f"  Service health: {health.get('status', 'unknown')}")

    # Generate encryption key
    index_key = Client.generate_key()

    # Create index
    index_name = f"wiki_full_index_test_{int(time.time())}"
    print(f"\n[4/8] Creating index '{index_name}'...")

    index_config = IndexIVFFlat(dimension=vector_dim)
    index = client.create_index(
        index_name=index_name,
        index_key=index_key,
        index_config=index_config,
        metric="euclidean",
    )
    print("  Index created successfully")

    try:
        # Insert vectors in batches
        print(f"\n[5/8] Inserting {actual_vector_count:,} vectors...")
        insert_start = time.time()
        total_inserted = 0

        for batch_start in tqdm(
            range(0, actual_vector_count, batch_size), desc="Upserting"
        ):
            batch_end = min(batch_start + batch_size, actual_vector_count)

            # Use integer IDs to match ground truth format
            ids = [str(i) for i in range(batch_start, batch_end)]
            vectors = train_vectors[batch_start:batch_end]

            index.upsert(ids, vectors)
            total_inserted += len(ids)

        insert_time = time.time() - insert_start
        insert_rate = total_inserted / insert_time

        print("\nInsert Results:")
        print(f"  Total vectors: {total_inserted:,}")
        print(f"  Total time: {insert_time:.2f}s")
        print(f"  Throughput: {insert_rate:,.2f} vectors/s")

        # Train the index (auto training)
        print("\n[6/8] Training index (auto)...")
        train_start = time.time()
        index.train(n_lists=0)

        # Wait for training to complete (it may be async)
        print("  Waiting for training to complete...")
        max_wait = 600  # 10 minutes max
        wait_start = time.time()
        while index.is_training():
            if time.time() - wait_start > max_wait:
                print("  WARNING: Training is taking longer than expected")
                break
            time.sleep(2)
            elapsed = time.time() - wait_start
            print(f"  Still training... ({elapsed:.0f}s)")

        if not index.is_trained():
            raise RuntimeError("Index training failed")

        train_time = time.time() - train_start
        print(f"  Training completed in {train_time:.2f}s")

        # Use auto n_probes (None) - the service will calculate optimal value
        n_probes = None
        print("  Using n_probes: auto")

        # Run queries and measure recall
        print(f"\n[7/8] Running {num_queries} queries...")
        query_times = []
        recalls = []

        for i in tqdm(range(num_queries), desc="Querying"):
            query_vector = query_vectors[i]

            query_start = time.time()
            results = index.query(
                query_vectors=query_vector, top_k=top_k, n_probes=n_probes
            )
            query_time = time.time() - query_start
            query_times.append(query_time)

            # Extract result IDs and compute recall against precomputed ground truth
            if results and len(results) > 0:
                result_ids = [r.get("id") or r.id for r in results]
                recall = compute_recall(ground_truth[i], result_ids)
                recalls.append(recall)

                if i == 0:
                    print(f"  First query returned {len(results)} results")
            else:
                if i == 0:
                    print("  WARNING: First query returned no results")
                recalls.append(0.0)

        avg_query_time = np.mean(query_times)
        p50_query_time = np.percentile(query_times, 50)
        p95_query_time = np.percentile(query_times, 95)
        p99_query_time = np.percentile(query_times, 99)

        print("\nQuery Results:")
        print(f"  Total queries: {num_queries}")
        print(f"  Average latency: {avg_query_time * 1000:.2f}ms")
        print(f"  P50 latency: {p50_query_time * 1000:.2f}ms")
        print(f"  P95 latency: {p95_query_time * 1000:.2f}ms")
        print(f"  P99 latency: {p99_query_time * 1000:.2f}ms")
        print(f"  QPS: {1 / avg_query_time:.2f}")

        # Recall results
        avg_recall = np.mean(recalls)
        min_recall = np.min(recalls)
        max_recall = np.max(recalls)

        print(f"\nRecall Results (Recall@{top_k}):")
        print(f"  Average recall: {avg_recall:.2%}")
        print(f"  Min recall: {min_recall:.2%}")
        print(f"  Max recall: {max_recall:.2%}")
        print(f"  Baseline threshold: {recall_baseline:.2%}")

        if avg_recall >= recall_baseline:
            print(
                f"  ✓ PASS: Recall {avg_recall:.2%} >= baseline {recall_baseline:.2%}"
            )
            test_passed = True
        else:
            print(f"  ✗ FAIL: Recall {avg_recall:.2%} < baseline {recall_baseline:.2%}")
            test_passed = False

        # Summary
        print("\n" + "=" * 60)
        print("FULL INDEX TEST COMPLETED")
        print("=" * 60)
        print("Dataset: wiki_all_1M")
        print(f"Index: {index_name}")
        print(f"Vectors indexed: {total_inserted:,}")
        print(f"Vector dimension: {vector_dim}")
        print(f"Insert throughput: {insert_rate:,.2f} vectors/s")
        print(f"Training time: {train_time:.2f}s")
        print(f"Query latency (avg): {avg_query_time * 1000:.2f}ms")
        print(f"Recall@{top_k}: {avg_recall:.2%}")
        print("=" * 60)

        if not test_passed:
            print("\nERROR: Recall below baseline threshold!")
            return 1

    finally:
        # Cleanup
        print("\n[8/8] Cleaning up...")
        try:
            index.delete_index()
            print(f"  Index '{index_name}' deleted")
        except Exception as e:
            print(f"  Warning: Failed to delete index: {e}")

    print("\nTest completed successfully!")
    return 0


if __name__ == "__main__":
    sys.exit(main())
