/**
 * Embedding generation using nomic-embed-text-v1.5 via @huggingface/transformers.
 *
 * Uses nomic's prefix convention:
 *   "search_document: " for content being indexed
 *   "search_query: "    for query strings at search time
 *
 * Model is lazy-loaded on first use and cached for the process lifetime.
 * ~274MB download on first call; stored in HF_CACHE_DIR.
 */

import { env, pipeline } from '@huggingface/transformers';
import path from 'path';

// Cache model files in a persistent location on the host
env.cacheDir = process.env.HF_CACHE_DIR ?? '/var/lib/nanoclaw/.cache/transformers';

/** Output dimension of nomic-embed-text-v1.5 */
export const EMBED_DIM = 768;

const MODEL_ID = 'nomic-ai/nomic-embed-text-v1.5';

type EmbedPipeline = Awaited<ReturnType<typeof pipeline>>;
let embedPipeline: EmbedPipeline | null = null;
let initPromise: Promise<void> | null = null;

async function initEmbedder(): Promise<void> {
  embedPipeline = await pipeline('feature-extraction', MODEL_ID, {
    dtype: 'fp32',
  });
}

async function getEmbedder(): Promise<EmbedPipeline> {
  if (!initPromise) {
    initPromise = initEmbedder();
  }
  await initPromise;
  return embedPipeline!;
}

/**
 * Generate a 768-dim embedding for the given text.
 * @param text  The text to embed (turn content or search query)
 * @param isQuery  true → prepend "search_query: ", false → "search_document: "
 */
export async function embedText(text: string, isQuery = false): Promise<Float32Array> {
  const prefix = isQuery ? 'search_query: ' : 'search_document: ';
  const embedder = await getEmbedder();
  const output = await (embedder as (text: string, options: Record<string, unknown>) => Promise<{ data: Float32Array }>)(
    prefix + text,
    { pooling: 'mean', normalize: true },
  );
  return output.data;
}

/**
 * Cosine similarity between two equal-length Float32Arrays.
 * Both vectors should be unit-normalised (nomic does this by default).
 * Returns a value in [-1, 1]; higher = more similar.
 */
export function cosineSim(a: Float32Array, b: Float32Array): number {
  let dot = 0;
  for (let i = 0; i < a.length; i++) dot += a[i] * b[i];
  // Since nomic normalises embeddings, magnitudes are 1 — dot product == cosine similarity.
  return dot;
}

/**
 * Pre-warm the embedding model at startup.
 * Fire-and-forget — call once in main() when LCM is enabled.
 */
export function warmupEmbedder(): void {
  if (!initPromise) {
    initPromise = initEmbedder().catch((err) => {
      // Reset so a retry is possible on next call
      initPromise = null;
      throw err;
    });
  }
}
