// EmbeddingsStore.ts
import fs from 'fs';
import path from 'path';

export type EmbeddingRow = {
  _id: string | number;
  embedding: number[];
  norm: number;
};

export type EmbeddingEntry = {
  v: number[];
  n: number;
};

export type EmbeddingsMap = Map<string, EmbeddingEntry>;
export type EmbeddingsByCollection = Record<string, EmbeddingsMap>;

export type WatchOptions = {
  debounceMs?: number;
};

export class EmbeddingsStore {
  private embDir: string;
  private active: string[];
  private maps: EmbeddingsByCollection;
  private watcher: fs.FSWatcher | null;

  constructor(embDir: string) {
    this.embDir = embDir;
    this.active = [];
    this.maps = {};
    this.watcher = null;
    this.refresh();
  }

  listActiveCollections = (): string[] => {
    if (!fs.existsSync(this.embDir)) return [];
    return fs
      .readdirSync(this.embDir)
      .filter((f) => f.endsWith('.jsonl'))
      .map((f) => path.basename(f, '.jsonl'));
  };

  loadEmbeddings = (collectionName: string): EmbeddingsMap => {
    const file = path.join(this.embDir, `${collectionName}.jsonl`);
    const map: EmbeddingsMap = new Map();
    if (!fs.existsSync(file)) return map;

    const raw = fs.readFileSync(file, 'utf8').trim();
    if (!raw) return map;

    for (const line of raw.split('\n')) {
      if (!line) continue;
      const r = JSON.parse(line) as EmbeddingRow;
      map.set(String(r._id), { v: r.embedding, n: r.norm });
    }

    return map;
  };

  refresh = (): this => {
    this.active = this.listActiveCollections();
    this.maps = Object.fromEntries(
      this.active.map((name) => [name, this.loadEmbeddings(name)]),
    );
    return this;
  };

  getActive = (): string[] => this.active;

  getMap = (collectionName: string): EmbeddingsMap =>
    this.maps[collectionName] ?? new Map<string, EmbeddingEntry>();

  startWatching = (opts: WatchOptions = {}): void => {
    const debounceMs = opts.debounceMs ?? 200;
    if (this.watcher) return;

    let t: NodeJS.Timeout | null = null;
    const schedule = () => {
      if (t) clearTimeout(t);
      t = setTimeout(() => this.refresh(), debounceMs);
    };

    if (!fs.existsSync(this.embDir)) return;

    this.watcher = fs.watch(
      this.embDir,
      { persistent: false },
      (_eventType, filename: any) => {
        // filename can be string | Buffer depending on Node version/options
        const name =
          typeof filename === 'string' ? filename : filename?.toString();
        if (!name) return;
        if (!name.endsWith('.jsonl')) return;
        schedule();
      },
    );
  };

  stopWatching = (): void => {
    if (!this.watcher) return;
    this.watcher.close();
    this.watcher = null;
  };
}