/**
 * Persistence Test for CyborgDB TypeScript SDK
 *
 * This script tests data persistence across service restarts:
 * - Phase 1 (setup): Creates index, inserts data, saves index info to file
 * - Phase 2 (verify): Loads index, verifies data still exists
 *
 * Usage:
 *   npx ts-node persistence_test_ts.ts setup    # Run before restart
 *   npx ts-node persistence_test_ts.ts verify   # Run after restart
 */

// 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';

const STATE_FILE = '/tmp/persistence_test_state_ts.json';
const VECTOR_COUNT = 1000;
const VECTOR_DIM = 128;

interface TestState {
  indexName: string;
  indexKey: string; // base64 encoded
  vectorCount: number;
  vectorDim: number;
  testIds: string[];
}

// Seeded random number generator for reproducible vectors
function seededRandom(seed: number): () => number {
  return function() {
    seed = (seed * 1103515245 + 12345) & 0x7fffffff;
    return (seed / 0x7fffffff) - 0.5;
  };
}

function generateSeededVector(dim: number, seed: number): number[] {
  const rng = seededRandom(seed);
  return Array.from({ length: dim }, () => rng() * 2);
}

function uint8ArrayToBase64(arr: Uint8Array): string {
  return Buffer.from(arr).toString('base64');
}

function base64ToUint8Array(base64: string): Uint8Array {
  return new Uint8Array(Buffer.from(base64, 'base64'));
}

async function setupPhase(): Promise<number> {
  console.log('============================================================');
  console.log('Persistence Test - Setup Phase');
  console.log('============================================================');

  const baseUrl = process.env.CYBORGDB_URL || 'http://localhost:8000';
  const apiKey = process.env.CYBORGDB_API_KEY;

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

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

  // Generate key and create index
  console.log('[2/4] Creating index...');
  const indexKey = client.generateKey();
  const indexName = 'persistence_test_index_ts';

  // Try to delete existing index
  try {
    const existing = await client.loadIndex({ indexName, indexKey });
    await existing.deleteIndex();
    console.log(`  Deleted existing index '${indexName}'`);
  } catch {
    // Index doesn't exist, that's fine
  }

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

  const index = await client.createIndex({
    indexName,
    indexKey,
    indexConfig,
    metric: 'euclidean',
  });
  console.log(`  Created index '${indexName}'`);

  // Insert test vectors with seeded random values
  console.log(`[3/4] Inserting ${VECTOR_COUNT} test vectors...`);

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

  for (let i = 0; i < VECTOR_COUNT; i++) {
    const id = `persist_vec_${i}`;
    ids.push(id);
    vectors.push(generateSeededVector(VECTOR_DIM, 42 + i));
    if (i < 10) {
      testIds.push(id);
    }
  }

  await index.upsert({ ids, vectors });
  console.log(`  Inserted ${VECTOR_COUNT} vectors`);

  // Verify data was inserted
  const listResult = await index.listIds();
  console.log(`  Verified ${listResult.ids.length} vectors in index`);

  // Save state for verification phase
  console.log('[4/4] Saving state...');
  const state: TestState = {
    indexName,
    indexKey: uint8ArrayToBase64(indexKey),
    vectorCount: VECTOR_COUNT,
    vectorDim: VECTOR_DIM,
    testIds,
  };

  fs.writeFileSync(STATE_FILE, JSON.stringify(state, null, 2));
  console.log(`  State saved to ${STATE_FILE}`);

  console.log('\n============================================================');
  console.log('Setup phase completed successfully!');
  console.log('Service can now be restarted.');
  console.log('============================================================');
  return 0;
}

async function verifyPhase(): Promise<number> {
  console.log('============================================================');
  console.log('Persistence Test - Verify Phase');
  console.log('============================================================');

  const baseUrl = process.env.CYBORGDB_URL || 'http://localhost:8000';
  const apiKey = process.env.CYBORGDB_API_KEY;

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

  // Load state from setup phase
  console.log('\n[1/4] Loading state...');
  if (!fs.existsSync(STATE_FILE)) {
    console.error(`ERROR: State file not found: ${STATE_FILE}`);
    console.log('Did you run the setup phase first?');
    return 1;
  }

  const state: TestState = JSON.parse(fs.readFileSync(STATE_FILE, 'utf-8'));
  const indexKey = base64ToUint8Array(state.indexKey);

  console.log(`  Index: ${state.indexName}`);
  console.log(`  Expected vectors: ${state.vectorCount}`);

  // Initialize client
  console.log('[2/4] Initializing client...');
  const client = new Client({ baseUrl, apiKey });

  // Load existing index
  console.log('[3/4] Loading index...');
  let index;
  try {
    index = await client.loadIndex({
      indexName: state.indexName,
      indexKey,
    });
    console.log(`  Index '${state.indexName}' loaded successfully`);
  } catch (e) {
    console.error(`ERROR: Failed to load index: ${e}`);
    return 1;
  }

  // Verify data
  console.log('[4/4] Verifying data...');

  // Check vector count
  const listResult = await index.listIds();
  const actualCount = listResult.ids.length;
  console.log(`  Found ${actualCount} vectors (expected ${state.vectorCount})`);

  if (actualCount !== state.vectorCount) {
    console.error('ERROR: Vector count mismatch!');
    return 1;
  }

  // Verify specific vectors exist
  try {
    const items = await index.get({ ids: state.testIds, include: ['vector'] });
    if (items.length !== state.testIds.length) {
      console.error(`ERROR: Expected ${state.testIds.length} items, got ${items.length}`);
      return 1;
    }
    console.log(`  Verified ${state.testIds.length} test vectors exist`);
  } catch (e) {
    console.error(`ERROR: Failed to get test vectors: ${e}`);
    return 1;
  }

  // Test query still works
  const queryVector = generateSeededVector(state.vectorDim, 42);
  const results = await index.query({
    queryVectors: queryVector,
    topK: 10,
  });

  if (!results.results || (Array.isArray(results.results) && results.results.length === 0)) {
    console.error('ERROR: Query returned no results');
    return 1;
  }
  console.log('  Query returned results successfully');

  // Cleanup
  console.log('\nCleaning up...');
  try {
    await index.deleteIndex();
    fs.unlinkSync(STATE_FILE);
    console.log('  Deleted index and state file');
  } catch (e) {
    console.log(`Warning: Cleanup failed: ${e}`);
  }

  console.log('\n============================================================');
  console.log('PERSISTENCE TEST PASSED!');
  console.log('Data successfully persisted across service restart.');
  console.log('============================================================');
  return 0;
}

async function main(): Promise<number> {
  if (process.argv.length < 3) {
    console.log('Usage: npx ts-node persistence_test_ts.ts <setup|verify>');
    return 1;
  }

  const phase = process.argv[2].toLowerCase();

  switch (phase) {
    case 'setup':
      return setupPhase();
    case 'verify':
      return verifyPhase();
    default:
      console.log(`Unknown phase: ${phase}`);
      console.log('Usage: npx ts-node persistence_test_ts.ts <setup|verify>');
      return 1;
  }
}

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