/**
 * Full Index Test for CyborgDB TypeScript 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
 *
 * Usage: npx ts-node full_index_test_ts.ts
 */

// Configure undici (Node.js fetch) timeouts for long-running operations like training
// Must be done before any fetch calls
// We need to replace globalThis.fetch because Node.js's built-in fetch doesn't respect setGlobalDispatcher
import { Agent, fetch as undiciFetch, setGlobalDispatcher } from 'undici';
const agent = new Agent({
  headersTimeout: 1800000,   // 30 minutes for headers (training 1M vectors takes a while)
  bodyTimeout: 1800000,      // 30 minutes for body
  connectTimeout: 30000,     // 30 seconds to connect
});
setGlobalDispatcher(agent);
// Replace global fetch with undici's fetch that respects our custom agent
(globalThis as any).fetch = undiciFetch;

// Import from built SDK dist folder (test runs from sdk directory after npm run build)
import { Client, IndexIVFFlat } from './dist';
import * as fs from 'fs';
import * as path from 'path';
import * as https from 'https';
// h5wasm is ESM-only, so we use dynamic import in loadDataset()

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

// Helper functions
function getEnvInt(key: string, defaultVal: number): number {
  const val = process.env[key];
  if (val) {
    const parsed = parseInt(val, 10);
    if (!isNaN(parsed)) return parsed;
  }
  return defaultVal;
}

function getEnvFloat(key: string, defaultVal: number): number {
  const val = process.env[key];
  if (val) {
    const parsed = parseFloat(val);
    if (!isNaN(parsed)) return parsed;
  }
  return defaultVal;
}

function percentile(data: number[], p: number): number {
  if (data.length === 0) return 0;
  const sorted = [...data].sort((a, b) => a - b);
  const idx = Math.floor((sorted.length - 1) * p / 100);
  return sorted[idx];
}

// Download dataset if not cached
async function downloadDataset(cacheDir: string): Promise<string> {
  if (!fs.existsSync(cacheDir)) {
    fs.mkdirSync(cacheDir, { recursive: true });
  }

  const datasetPath = path.join(cacheDir, DATASET_FILENAME);

  if (fs.existsSync(datasetPath)) {
    console.log(`  Dataset already cached at ${datasetPath}`);
    return datasetPath;
  }

  console.log(`  Downloading from ${DATASET_URL}...`);

  return new Promise((resolve, reject) => {
    const file = fs.createWriteStream(datasetPath);
    https.get(DATASET_URL, (response) => {
      if (response.statusCode === 302 || response.statusCode === 301) {
        // Handle redirect
        https.get(response.headers.location!, (redirectResponse) => {
          redirectResponse.pipe(file);
          file.on('finish', () => {
            file.close();
            console.log(`  Downloaded to ${datasetPath}`);
            resolve(datasetPath);
          });
        }).on('error', reject);
      } else {
        response.pipe(file);
        file.on('finish', () => {
          file.close();
          console.log(`  Downloaded to ${datasetPath}`);
          resolve(datasetPath);
        });
      }
    }).on('error', (err) => {
      fs.unlink(datasetPath, () => {});
      reject(err);
    });
  });
}

// Load dataset from HDF5 file
async function loadDataset(
  datasetPath: string,
  numVectors: number,
  numQueries: number,
  topK: number
): Promise<{
  trainVectors: number[][];
  queryVectors: number[][];
  groundTruth: number[][];
  vectorDim: number;
}> {
  // Use h5wasm/node which has NODERAWFS enabled for direct filesystem access
  // This allows reading large files (>2GB) without loading into memory
  const h5wasm = await import('h5wasm/node');
  await h5wasm.ready;

  const absolutePath = path.resolve(datasetPath);

  // Check file exists and get size
  const stats = fs.statSync(absolutePath);
  console.log(`  HDF5 file: ${absolutePath}`);
  console.log(`  File size: ${(stats.size / 1024 / 1024 / 1024).toFixed(2)} GB`);

  // With h5wasm/node (NODERAWFS), we can directly open files using absolute paths
  console.log(`  Opening HDF5 file with NODERAWFS...`);
  const file = new h5wasm.File(absolutePath, 'r');

  // Load train vectors - use chunked slice to avoid WASM memory limits (2GB max)
  // 1M vectors * 768 dims * 4 bytes = ~3GB, so we read in 100K chunks
  const trainDataset = file.get('train') as any;
  const trainShape = trainDataset.shape as number[];
  const totalVectors = trainShape[0];
  const vectorDim = trainShape[1];

  if (numVectors > totalVectors) {
    numVectors = totalVectors;
  }

  const CHUNK_SIZE = 100000;  // 100K vectors per chunk to stay under WASM 2GB limit
  console.log(`  Reading ${numVectors.toLocaleString()} train vectors in chunks of ${CHUNK_SIZE.toLocaleString()}...`);
  const trainVectors: number[][] = [];

  for (let start = 0; start < numVectors; start += CHUNK_SIZE) {
    const end = Math.min(start + CHUNK_SIZE, numVectors);
    const chunkData = trainDataset.slice([[start, end]]) as Float32Array;

    // Convert chunk to vectors
    for (let i = 0; i < (end - start); i++) {
      trainVectors.push(Array.from(chunkData.slice(i * vectorDim, (i + 1) * vectorDim)));
    }

    const progress = ((trainVectors.length / numVectors) * 100).toFixed(0);
    process.stdout.write(`\r  Loading train vectors: ${trainVectors.length.toLocaleString()}/${numVectors.toLocaleString()} (${progress}%)`);
  }
  console.log();  // New line after progress

  // Load test vectors (queries) - use slice
  const testDataset = file.get('test') as any;
  const testShape = testDataset.shape as number[];
  const totalQueries = testShape[0];

  if (numQueries > totalQueries) {
    numQueries = totalQueries;
  }

  console.log(`  Reading ${numQueries.toLocaleString()} test vectors using slice...`);
  const testData = testDataset.slice([[0, numQueries]]) as Float32Array;
  const queryVectors: number[][] = [];
  for (let i = 0; i < numQueries; i++) {
    queryVectors.push(Array.from(testData.slice(i * vectorDim, (i + 1) * vectorDim)));
  }

  // Load ground truth neighbors - use slice
  const neighborsDataset = file.get('neighbors') as any;
  const neighborsShape = neighborsDataset.shape as number[];
  const neighborsK = neighborsShape[1];

  if (topK > neighborsK) {
    topK = neighborsK;
  }

  console.log(`  Reading ${numQueries.toLocaleString()} neighbor rows using slice...`);
  const neighborsData = neighborsDataset.slice([[0, numQueries]]) as Int32Array;
  const groundTruth: number[][] = [];
  for (let i = 0; i < numQueries; i++) {
    groundTruth.push(Array.from(neighborsData.slice(i * neighborsK, i * neighborsK + topK)));
  }

  file.close();

  return { trainVectors, queryVectors, groundTruth, vectorDim };
}

// Compute recall between ground truth and ANN results
function computeRecall(groundTruth: number[], resultIds: string[]): number {
  if (groundTruth.length === 0) return 0;

  const gtSet = new Set(groundTruth);
  let matches = 0;
  for (const idStr of resultIds) {
    const id = parseInt(idStr, 10);
    if (gtSet.has(id)) matches++;
  }
  return matches / groundTruth.length;
}

async function main(): Promise<number> {
  // Configuration from environment
  const baseUrl = process.env.CYBORGDB_URL || 'http://localhost:8000';
  const apiKey = process.env.CYBORGDB_API_KEY;

  if (!apiKey) {
    console.error('ERROR: CYBORGDB_API_KEY environment variable not set');
    return 1;
  }

  const vectorCount = getEnvInt('VECTOR_COUNT', 1000000);
  const batchSize = getEnvInt('BATCH_SIZE', 10000);
  const topK = getEnvInt('TOP_K', 100);
  const numQueries = getEnvInt('NUM_QUERIES', 1000);
  const recallBaseline = getEnvFloat('RECALL_BASELINE', 0.95);
  const cacheDir = process.env.DATASET_CACHE_DIR || '/tmp/ann-benchmarks';

  console.log('============================================================');
  console.log('CyborgDB Full Index Test - TypeScript SDK (wiki_all_1M)');
  console.log('============================================================');
  console.log(`Service URL: ${baseUrl}`);
  console.log(`Vector count: ${vectorCount.toLocaleString()}`);
  console.log(`Batch size: ${batchSize.toLocaleString()}`);
  console.log(`Top K: ${topK}`);
  console.log(`Number of queries: ${numQueries}`);
  console.log(`Recall baseline: ${(recallBaseline * 100).toFixed(0)}%`);
  console.log('============================================================');

  // Download dataset
  console.log('\n[1/8] Downloading/loading wiki_all_1M dataset...');
  const datasetPath = await downloadDataset(cacheDir);

  // Load dataset
  console.log('\n[2/8] Loading dataset...');
  const { trainVectors, queryVectors, groundTruth, vectorDim } = await loadDataset(
    datasetPath,
    vectorCount,
    numQueries,
    topK
  );
  const actualVectorCount = trainVectors.length;

  console.log(`  Train vectors: ${actualVectorCount} x ${vectorDim}`);
  console.log(`  Query vectors: ${queryVectors.length} x ${vectorDim}`);
  console.log(`  Ground truth: ${groundTruth.length} x ${groundTruth[0].length}`);

  // Initialize client
  console.log('\n[3/8] Initializing client...');
  const client = new Client({
    baseUrl,
    apiKey,
  });

  // Check health
  try {
    const health = await client.getHealth();
    console.log(`  Service health: ${health.status || 'unknown'}`);
  } catch (e) {
    console.log(`  WARNING: Health check failed: ${e}`);
  }

  // Generate encryption key
  const indexKey = client.generateKey();

  // Create index
  const indexName = `wiki_full_index_test_${Date.now()}`;
  console.log(`\n[4/8] Creating index '${indexName}'...`);

  const indexConfig: IndexIVFFlat = {
    dimension: vectorDim,
    type: 'ivfflat',
  };

  const index = await client.createIndex({
    indexName,
    indexKey,
    indexConfig,
    metric: 'euclidean',
  });
  console.log('  Index created successfully');

  // Track test result
  let testPassed = true;

  try {
    // Insert vectors in batches
    console.log(`\n[5/8] Inserting ${actualVectorCount.toLocaleString()} vectors...`);
    const insertStart = Date.now();
    let totalInserted = 0;

    for (let batchStart = 0; batchStart < actualVectorCount; batchStart += batchSize) {
      const batchEnd = Math.min(batchStart + batchSize, actualVectorCount);
      const currentBatchSize = batchEnd - batchStart;

      const ids: string[] = [];
      const vectors: number[][] = [];

      for (let i = 0; i < currentBatchSize; i++) {
        ids.push(String(batchStart + i));
        vectors.push(trainVectors[batchStart + i]);
      }

      await index.upsert({ ids, vectors });
      totalInserted += currentBatchSize;

      const progress = ((totalInserted / actualVectorCount) * 100).toFixed(1);
      process.stdout.write(`\r  Progress: ${totalInserted}/${actualVectorCount} (${progress}%)`);
    }

    const insertTime = (Date.now() - insertStart) / 1000;
    const insertRate = totalInserted / insertTime;

    console.log(`\n\nInsert Results:`);
    console.log(`  Total vectors: ${totalInserted.toLocaleString()}`);
    console.log(`  Total time: ${insertTime.toFixed(2)}s`);
    console.log(`  Throughput: ${insertRate.toFixed(2)} vectors/s`);

    // Train the index (auto training)
    console.log(`\n[6/8] Training index (auto)...`);
    const trainStart = Date.now();
    await index.train({ nLists: 0 });
    const trainTime = (Date.now() - trainStart) / 1000;
    console.log(`  Training completed in ${trainTime.toFixed(2)}s`);

    // Use auto n_probes (undefined) - the service will calculate optimal value
    const nProbes = undefined;
    console.log(`  Using n_probes: auto`);

    // Run queries and measure recall
    console.log(`\n[7/8] Running ${numQueries} queries...`);
    const queryTimes: number[] = [];
    const recalls: number[] = [];

    for (let i = 0; i < numQueries; i++) {
      const queryVector = queryVectors[i];

      const queryStart = Date.now();
      const results = await index.query({
        queryVectors: queryVector,
        topK,
        nProbes,
      });
      const queryTime = (Date.now() - queryStart) / 1000;
      queryTimes.push(queryTime);

      // Extract result IDs and compute recall
      if (results.results && Array.isArray(results.results) && results.results.length > 0) {
        const resultIds = results.results.map((r: any) => r.id || r.Id);
        const recall = computeRecall(groundTruth[i], resultIds);
        recalls.push(recall);

        if (i === 0) {
          console.log(`  First query returned ${results.results.length} results`);
        }
      } else {
        if (i === 0) {
          console.log('  WARNING: First query returned no results');
        }
        recalls.push(0.0);
      }

      if ((i + 1) % 100 === 0) {
        process.stdout.write(`\r  Progress: ${i + 1}/${numQueries}`);
      }
    }

    // Calculate query statistics
    const avgQueryTime = queryTimes.reduce((a, b) => a + b, 0) / numQueries;
    const p50QueryTime = percentile(queryTimes, 50);
    const p95QueryTime = percentile(queryTimes, 95);
    const p99QueryTime = percentile(queryTimes, 99);

    console.log(`\n\nQuery Results:`);
    console.log(`  Total queries: ${numQueries}`);
    console.log(`  Average latency: ${(avgQueryTime * 1000).toFixed(2)}ms`);
    console.log(`  P50 latency: ${(p50QueryTime * 1000).toFixed(2)}ms`);
    console.log(`  P95 latency: ${(p95QueryTime * 1000).toFixed(2)}ms`);
    console.log(`  P99 latency: ${(p99QueryTime * 1000).toFixed(2)}ms`);
    console.log(`  QPS: ${(1 / avgQueryTime).toFixed(2)}`);

    // Calculate recall statistics
    const avgRecall = recalls.reduce((a, b) => a + b, 0) / recalls.length;
    const minRecall = Math.min(...recalls);
    const maxRecall = Math.max(...recalls);

    console.log(`\nRecall Results (Recall@${topK}):`);
    console.log(`  Average recall: ${(avgRecall * 100).toFixed(2)}%`);
    console.log(`  Min recall: ${(minRecall * 100).toFixed(2)}%`);
    console.log(`  Max recall: ${(maxRecall * 100).toFixed(2)}%`);
    console.log(`  Baseline threshold: ${(recallBaseline * 100).toFixed(2)}%`);

    if (avgRecall >= recallBaseline) {
      console.log(`  ✓ PASS: Recall ${(avgRecall * 100).toFixed(2)}% >= baseline ${(recallBaseline * 100).toFixed(2)}%`);
    } else {
      console.log(`  ✗ FAIL: Recall ${(avgRecall * 100).toFixed(2)}% < baseline ${(recallBaseline * 100).toFixed(2)}%`);
      testPassed = false;
    }

    // Summary
    console.log('\n============================================================');
    console.log('FULL INDEX TEST COMPLETED');
    console.log('============================================================');
    console.log(`Dataset: wiki_all_1M`);
    console.log(`Index: ${indexName}`);
    console.log(`Vectors indexed: ${totalInserted.toLocaleString()}`);
    console.log(`Vector dimension: ${vectorDim}`);
    console.log(`Insert throughput: ${insertRate.toFixed(2)} vectors/s`);
    console.log(`Training time: ${trainTime.toFixed(2)}s`);
    console.log(`Query latency (avg): ${(avgQueryTime * 1000).toFixed(2)}ms`);
    console.log(`Recall@${topK}: ${(avgRecall * 100).toFixed(2)}%`);
    console.log('============================================================');

  } finally {
    // Cleanup
    console.log('\n[8/8] Cleaning up...');
    try {
      await index.deleteIndex();
      console.log(`  Index '${indexName}' deleted`);
    } catch (e) {
      console.log(`  Warning: Failed to delete index: ${e}`);
    }
  }

  if (testPassed) {
    console.log('\nTest completed successfully!');
    return 0;
  } else {
    console.log('\nTest FAILED: Recall below baseline threshold!');
    return 1;
  }
}

// Run main
main()
  .then((code) => process.exit(code))
  .catch((err) => {
    console.error('ERROR:', err);
    process.exit(1);
  });
