// Drives on-device indexing for semantic search (Phase 1): loads the embedding
// model, embeds each transcript's segments into unit vectors, and persists them
// to the StorageRepo's vector store. UI/session state only — the repo is the
// record of truth for what's been embedded.
import { create } from 'zustand';

import { getRepo } from '@/lib/db';
import { EMBED_MODEL, getEmbeddingEngine } from '@/lib/embedding';

type Status = 'idle' | 'indexing' | 'ready';

interface EmbeddingState {
  status: Status;
  /** Build progress in [0,1] (0.2 = model load, 0.8 = embedding). */
  progress: number;
  /** Count of transcripts with no vectors for the active model. */
  pending: number;

  /** Refresh `pending` from the repo (cheap; no model load). */
  refreshPending: () => Promise<void>;
  /** Embed every un-embedded transcript. Loads the model first. */
  buildIndex: () => Promise<void>;
  /** Embed a single transcript (used right after a new transcription). */
  embedOne: (transcriptId: string) => Promise<void>;
}

/**
 * Embed and persist vectors for one transcript. Replaces any existing vectors
 * for it (upsertVectors semantics). No-op if the transcript or its segments are
 * gone. Caller is responsible for model loading and status/progress bookkeeping.
 */
async function embedTranscript(transcriptId: string): Promise<void> {
  const repo = getRepo();
  const t = await repo.get(transcriptId);
  if (!t) return;
  if (t.segments.length === 0) {
    // Nothing to embed, but record an (empty) vector set so it stops counting
    // as "unembedded" for this model.
    await repo.upsertVectors(transcriptId, EMBED_MODEL, []);
    return;
  }
  const eng = getEmbeddingEngine();
  const vecs = await eng.embed(
    t.segments.map((s) => s.text),
    'passage',
  );
  await repo.upsertVectors(
    transcriptId,
    EMBED_MODEL,
    t.segments.map((s, j) => ({
      segmentId: s.id ?? String(j),
      start: s.start,
      end: s.end,
      text: s.text,
      vector: vecs[j] ?? new Float32Array(eng.dim),
    })),
  );
}

export const useEmbedding = create<EmbeddingState>((set, get) => ({
  status: 'idle',
  progress: 0,
  pending: 0,

  refreshPending: async () => {
    const ids = await getRepo().unembeddedIds(EMBED_MODEL);
    set({ pending: ids.length });
  },

  buildIndex: async () => {
    // Guard against concurrent runs.
    if (get().status === 'indexing') return;
    set({ status: 'indexing', progress: 0 });
    try {
      const eng = getEmbeddingEngine();
      // Model load is the first 20% of the bar.
      await eng.loadModel((p) => set({ progress: p * 0.2 }));

      const ids = await getRepo().unembeddedIds(EMBED_MODEL);
      if (ids.length === 0) {
        set({ status: 'ready', progress: 1, pending: 0 });
        return;
      }

      for (let i = 0; i < ids.length; i++) {
        const id = ids[i];
        if (id === undefined) continue;
        await embedTranscript(id);
        // Remaining 80% spread across the transcripts.
        set({ progress: 0.2 + 0.8 * ((i + 1) / ids.length) });
      }

      set({ status: 'ready', progress: 1, pending: 0 });
    } catch (err) {
      // Surface progress reset but don't crash callers; leave pending intact so
      // a retry can pick up where it left off.
      set({ status: 'idle', progress: 0 });
      throw err;
    }
  },

  embedOne: async (transcriptId) => {
    // Guard against clobbering an in-flight full build.
    if (get().status === 'indexing') return;
    set({ status: 'indexing' });
    try {
      const eng = getEmbeddingEngine();
      // Lazy model load (no progress weighting for a single transcript).
      if (!eng.isLoaded()) await eng.loadModel();
      await embedTranscript(transcriptId);
      const pending = (await getRepo().unembeddedIds(EMBED_MODEL)).length;
      set({ status: 'ready', pending });
    } catch (err) {
      set({ status: 'idle' });
      throw err;
    }
  },
}));
