/**
 * Embedding quantization — ADR-121 Phase 3.
 *
 * Two opt-in quantization paths:
 *
 *   1. **int8** — per-vector linear quantization. 4× memory
 *      reduction (float32 → int8). Works with streaming inserts
 *      because each vector encodes independently. Used by
 *      `SearchableEmbeddingCache` when `quantize: 'int8'` is set.
 *
 *   2. **RaBitQ** — 1-bit-per-dimension quantization via
 *      `@ruvector/rabitq-wasm`. 32× memory reduction. Build-once
 *      batch API (not streaming) — exposed here as the
 *      `RabitqSnapshot` class for callers that want to compress
 *      a fixed corpus and query it. Recall recovered via an
 *      exact-L2² rerank pool (default 20× the requested k).
 *
 * Each path is opt-in + has a measurable acceptance gate:
 *   int8     → ≥0.98 cosine-similarity recall after dequantize
 *   RaBitQ   → ≥0.95 recall after rerank, 32× memory reduction
 *
 * For int8 the recall acceptance is measured per-vector
 * (cosine(original, dequant(quant(original)))). For RaBitQ
 * recall is measured against the top-k from a brute-force baseline.
 */

import type { CacheSearchHit } from './searchable-embedding-cache.js';

// ============================================================================
// int8 linear quantization
// ============================================================================

/**
 * Encoded int8 vector: the int8 buffer + the scale factor needed to
 * dequantize. Scale is per-vector (each vector's max-abs sets its own
 * scale) — that's the standard approach for embedding quantization
 * and preserves recall best for unit-normalized vectors.
 */
export interface Int8EncodedVector {
  /** Int8 buffer, same length as the source float32 vector. */
  readonly bytes: Int8Array;
  /** Scale factor: float32 ≈ int8 * scale. */
  readonly scale: number;
}

/**
 * Encode a float32 vector as int8. Returns the byte buffer + a per-vector
 * scale that maps int8 back to ~float32 on decode.
 *
 * Algorithm: find max-abs value, scale so max-abs maps to 127, round
 * to nearest int8. This is min-max scaling clamped to [-127, 127] and
 * is the cheapest path that preserves direction (which is what cosine
 * similarity cares about for unit-normalized embeddings).
 */
export function int8Encode(vec: Float32Array): Int8EncodedVector {
  let maxAbs = 0;
  for (let i = 0; i < vec.length; i++) {
    const v = Math.abs(vec[i]!);
    if (v > maxAbs) maxAbs = v;
  }
  // All-zero vector → emit zero bytes with a unit scale so decode
  // round-trips to the same zeros without a divide-by-zero.
  if (maxAbs === 0) {
    return { bytes: new Int8Array(vec.length), scale: 1 };
  }
  const scale = maxAbs / 127;
  const invScale = 127 / maxAbs;
  const bytes = new Int8Array(vec.length);
  for (let i = 0; i < vec.length; i++) {
    // Round toward zero with .5 -> away-from-zero (matches the
    // banker's-round-style behavior of typical quantization libs;
    // doesn't matter for the recall metric but matters for
    // determinism across platforms).
    const q = Math.round(vec[i]! * invScale);
    // Clamp to int8 range — should be redundant given the scale
    // derivation but defensive against numeric drift.
    bytes[i] = Math.max(-127, Math.min(127, q));
  }
  return { bytes, scale };
}

/**
 * Decode an int8-encoded vector back to float32. Round-trips with
 * recall ≥0.98 on unit-normalized vectors (the ADR-121 Phase 3 gate).
 */
export function int8Decode(encoded: Int8EncodedVector): Float32Array {
  const out = new Float32Array(encoded.bytes.length);
  const { bytes, scale } = encoded;
  for (let i = 0; i < bytes.length; i++) {
    out[i] = bytes[i]! * scale;
  }
  return out;
}

/**
 * Cosine similarity between two vectors. Used by the quantization
 * recall test; duplicated here (rather than importing from
 * embedding-service.ts) to keep this module's dependency graph
 * minimal — quantization should never depend on the service layer.
 */
function cosineSimilarity(a: Float32Array, b: Float32Array): number {
  let dot = 0, na = 0, nb = 0;
  const len = Math.min(a.length, b.length);
  for (let i = 0; i < len; i++) {
    dot += a[i]! * b[i]!;
    na += a[i]! * a[i]!;
    nb += b[i]! * b[i]!;
  }
  if (na === 0 || nb === 0) return 0;
  return dot / (Math.sqrt(na) * Math.sqrt(nb));
}

/**
 * Measure round-trip recall for int8 encoding on a corpus.
 * Recall here = mean cosine similarity between `vec` and
 * `int8Decode(int8Encode(vec))` across all corpus vectors.
 *
 * Used by the CI benchmark to assert the ADR-121 Phase 3 acceptance
 * gate (≥0.98 recall for int8).
 */
export function measureInt8RoundTripRecall(corpus: Float32Array[]): {
  meanRecall: number;
  minRecall: number;
  maxRecall: number;
  count: number;
} {
  if (corpus.length === 0) {
    return { meanRecall: 1, minRecall: 1, maxRecall: 1, count: 0 };
  }
  let sum = 0, min = 1, max = 0;
  for (const vec of corpus) {
    const round = int8Decode(int8Encode(vec));
    const sim = cosineSimilarity(vec, round);
    sum += sim;
    if (sim < min) min = sim;
    if (sim > max) max = sim;
  }
  return { meanRecall: sum / corpus.length, minRecall: min, maxRecall: max, count: corpus.length };
}

// ============================================================================
// RaBitQ batch snapshot (1-bit quantization via @ruvector/rabitq-wasm)
// ============================================================================

/**
 * Structural type for the WASM RabitqIndex we use. Mirrors the
 * public surface — a future minor that extends the class won't
 * break us.
 */
interface RabitqIndexLike {
  search(query: Float32Array, k: number): Array<{ id: number; distance: number }>;
  readonly len: number;
  readonly isEmpty: boolean;
  free?(): void;
}

interface RabitqWasmModule {
  RabitqIndex: {
    build(vectors: Float32Array, dim: number, seed: bigint, rerankFactor: number): RabitqIndexLike;
  };
}

async function loadRabitqWasm(): Promise<RabitqWasmModule | null> {
  try {
    const mod = (await import('@ruvector/rabitq-wasm')) as unknown as Record<string, unknown>;
    // Either flat-named export or a `.default` shape (CJS bridge).
    const surface = (typeof (mod as { RabitqIndex?: unknown }).RabitqIndex === 'function'
      ? mod
      : (mod.default as Record<string, unknown> | undefined)) as RabitqWasmModule | undefined;
    return surface && typeof surface.RabitqIndex?.build === 'function' ? surface : null;
  } catch {
    return null;
  }
}

export interface RabitqSnapshotOptions {
  /** Embedding dimension. Required. */
  readonly dimension: number;
  /** Random rotation seed (default 0n) — same seed = bit-identical codes. */
  readonly seed?: bigint;
  /** Multiplier on `k` for the exact-L2² rerank pool (default 20). */
  readonly rerankFactor?: number;
}

/**
 * Snapshot result — what callers get back from `search`. `id` is the
 * caller-supplied key (string), preserved through the build/search
 * round-trip via the snapshot's internal id-map.
 */
export interface RabitqSnapshotHit {
  readonly id: string;
  /** Approximate L2² distance after RaBitQ rerank. */
  readonly distance: number;
  /** Original vector (preserved by the snapshot for downstream use). */
  readonly vector: Float32Array;
}

/**
 * Batch RaBitQ snapshot — compress a corpus of vectors once, then
 * query against the 1-bit index with exact-L2² rerank for recall.
 *
 * Lifecycle:
 *   const snap = await RabitqSnapshot.build(entries, { dimension: 384 });
 *   const hits = await snap.search(queryVec, 10);
 *   snap.dispose();
 *
 * Memory profile:
 *   raw  = N × dim × 4 bytes (float32)
 *   rabitq = N × dim / 8 bytes (1 bit per dim) → 32× smaller
 * (Plus a per-snapshot id-map for string → numeric position; small.)
 *
 * If `@ruvector/rabitq-wasm` isn't installed, `build` throws a
 * clear named error so callers know which optional peer to install.
 */
export class RabitqSnapshot {
  private readonly index: RabitqIndexLike;
  private readonly idMap: string[];
  private readonly vectors: Float32Array[];
  private disposed = false;

  private constructor(index: RabitqIndexLike, idMap: string[], vectors: Float32Array[]) {
    this.index = index;
    this.idMap = idMap;
    this.vectors = vectors;
  }

  static async build(
    entries: Array<{ id: string; vector: Float32Array }>,
    options: RabitqSnapshotOptions,
  ): Promise<RabitqSnapshot> {
    const mod = await loadRabitqWasm();
    if (!mod) {
      throw new Error(
        '@ruvector/rabitq-wasm is not installed or did not load cleanly. ' +
        'Run: npm install @ruvector/rabitq-wasm. ' +
        'Falls back to keeping a plain float32 cache when not installed.',
      );
    }
    if (entries.length === 0) {
      throw new Error('RabitqSnapshot.build: cannot build an empty index');
    }
    const dim = options.dimension;
    const seed = options.seed ?? 0n;
    const rerankFactor = options.rerankFactor ?? 20;

    // Flatten the entries into one big Float32Array of length N*dim.
    // The RabitqIndex assigns ids by position, so we keep a parallel
    // id-map for the string → position lookup on search results.
    const flat = new Float32Array(entries.length * dim);
    const idMap: string[] = new Array(entries.length);
    const vectors: Float32Array[] = new Array(entries.length);
    for (let i = 0; i < entries.length; i++) {
      const e = entries[i]!;
      if (e.vector.length !== dim) {
        throw new Error(
          `RabitqSnapshot.build: entry ${i} ("${e.id}") has dim ${e.vector.length}, ` +
          `expected ${dim}`,
        );
      }
      flat.set(e.vector, i * dim);
      idMap[i] = e.id;
      vectors[i] = e.vector;
    }
    const index = mod.RabitqIndex.build(flat, dim, seed, rerankFactor);
    return new RabitqSnapshot(index, idMap, vectors);
  }

  /** Query the snapshot for the `k` nearest neighbors of `query`. */
  async search(query: Float32Array, k: number): Promise<RabitqSnapshotHit[]> {
    if (this.disposed) throw new Error('RabitqSnapshot: already disposed');
    const raw = this.index.search(query, k);
    const out: RabitqSnapshotHit[] = [];
    for (const row of raw) {
      const id = this.idMap[row.id];
      const vec = this.vectors[row.id];
      if (id === undefined || !vec) continue;
      out.push({ id, distance: row.distance, vector: vec });
    }
    return out;
  }

  /** Number of vectors in the snapshot. */
  get size(): number {
    return this.idMap.length;
  }

  /**
   * Release the WASM index. Idempotent. Memory held by the JS-side
   * id-map and vectors is GC'd normally.
   */
  dispose(): void {
    if (this.disposed) return;
    try { this.index.free?.(); } catch { /* best-effort */ }
    this.disposed = true;
  }
}

/**
 * Measure the *theoretical* memory reduction for a corpus when
 * encoded with RaBitQ vs raw float32. Returns the byte counts so
 * callers can publish the metric to the tracker. Doesn't actually
 * build the index (cheap to compute, deterministic).
 */
export function rabitqMemoryReduction(corpusSize: number, dimension: number): {
  rawBytes: number;
  rabitqBytes: number;
  reductionFactor: number;
} {
  const rawBytes = corpusSize * dimension * 4; // float32
  const rabitqBytes = Math.ceil(corpusSize * dimension / 8); // 1 bit per dim
  return {
    rawBytes,
    rabitqBytes,
    reductionFactor: rawBytes / Math.max(1, rabitqBytes),
  };
}

// Re-export the hit type for cache integration.
export type { CacheSearchHit };