// Query embedder for the API: the EmbeddingGemma 2 *text* encoder at the same pinned revision and q4 weights as the
// app's WebGPU path (src/embed-worker.js), on onnxruntime-node's CPU backend. Vision/audio encoders are never loaded.
//   MODEL_DIR=<dir> node server/embed.mjs --prefetch
//       download the pinned files into <dir>/<model id>/ and load them exactly as the server will (Docker build step)
// With MODEL_DIR set, the model is loaded from that directory with remote access disabled. A plain local model
// directory is used rather than an offline revision-pinned cache because transformers.js 4.3.1 looks up the tokenizer
// files at the main revision (get_tokenizer_files ignores `revision`), which such a cache does not contain.
// Without MODEL_DIR (development), files come from the Hub at the pinned revision via the transformers.js cache.
import { AutoConfig, AutoModel, AutoTokenizer, env } from '@huggingface/transformers';
import fs from 'node:fs';
import path from 'node:path';
import { Readable } from 'node:stream';
import { pipeline } from 'node:stream/promises';
import { pathToFileURL } from 'node:url';

const MODEL_ID = 'onnx-community/embeddinggemma-2-ONNX';
const REVISION = 'daa72c51243991dfcaf9f9137d2c573d8f7790c0'; // same pin as src/embed-worker.js
const FILES = ['config.json', 'tokenizer.json', 'tokenizer_config.json', 'onnx/model_q4.onnx', 'onnx/model_q4.onnx_data'];
const QUERY_PREFIX = 'task: search result | query: ';
const MODEL_DIR = process.env.MODEL_DIR;
if (process.env.MODEL_CACHE) env.cacheDir = process.env.MODEL_CACHE;
if (MODEL_DIR) {
  env.localModelPath = path.resolve(MODEL_DIR);
  env.allowLocalModels = true;
  env.allowRemoteModels = false;
}

export async function loadEmbedder() {
  const t0 = performance.now();
  const pin = MODEL_DIR ? {} : { revision: REVISION }; // the local directory already holds the pinned revision
  const config = await AutoConfig.from_pretrained(MODEL_ID, pin);
  config.vision_config = null; // text-only, as in the browser
  config.audio_config = null;
  const [model, tokenizer] = await Promise.all([
    AutoModel.from_pretrained(MODEL_ID, { ...pin, config, device: 'cpu', dtype: 'q4' }),
    AutoTokenizer.from_pretrained(MODEL_ID, pin),
  ]);
  const cache = new Map(); // text -> Float32Array (small LRU)
  let chain = Promise.resolve(); // one inference at a time: requests queue instead of oversubscribing the CPU
  const embedNow = async (text) => {
    const inputs = await tokenizer([QUERY_PREFIX + text], { padding: true, truncation: true, max_length: 128 });
    const out = await model(inputs);
    return new Float32Array(out.sentence_embedding.data);
  };
  const s = {
    device: 'cpu', backend: 'onnxruntime-node', dtype: 'q4', model: `${MODEL_ID}@${REVISION.slice(0, 7)}`,
    lastMs: null,
    embed(text) {
      const hit = cache.get(text);
      if (hit) { cache.delete(text); cache.set(text, hit); return Promise.resolve(hit); }
      const run = chain.then(async () => {
        const t = performance.now();
        const v = await embedNow(text);
        s.lastMs = performance.now() - t;
        cache.set(text, v);
        if (cache.size > 500) cache.delete(cache.keys().next().value);
        return v;
      });
      chain = run.catch(() => {});
      return run;
    },
  };
  s.loadMs = performance.now() - t0;
  const w0 = performance.now();
  await embedNow('warm up');
  s.warmMs = performance.now() - w0;
  return s;
}

async function prefetch() {
  if (!MODEL_DIR) throw new Error('set MODEL_DIR');
  for (const file of FILES) {
    const dst = path.join(MODEL_DIR, MODEL_ID, file);
    const r = await fetch(`https://huggingface.co/${MODEL_ID}/resolve/${REVISION}/${file}`);
    if (!r.ok) throw new Error(`${r.status} fetching ${file}`);
    fs.mkdirSync(path.dirname(dst), { recursive: true });
    await pipeline(Readable.fromWeb(r.body), fs.createWriteStream(dst));
  }
}

if (import.meta.url === pathToFileURL(process.argv[1]).href && process.argv.includes('--prefetch')) {
  await prefetch();
  const s = await loadEmbedder(); // offline, from MODEL_DIR: fails the build if the server could not load it
  const v = await s.embed('a cat sleeping');
  console.log(JSON.stringify({ model: s.model, dir: path.join(MODEL_DIR, MODEL_ID), files: FILES.length, loadMs: Math.round(s.loadMs), warmMs: Math.round(s.warmMs), dims: v.length }));
}