import { AutoTokenizer, CLIPTextModelWithProjection, env } from '@huggingface/transformers';
import { cachedFile, DATASET, TOKENIZER, TEXT_MODEL, TEXT_MODEL_REVISION, DIMENSIONS } from './assets.js';
import { normalizeDatabase, rankImages } from './search.js';
import { loadCatalogue, searchCatalogue, tryLoadCatalogue } from './catalogue.js';

env.allowLocalModels = false;
env.backends.onnx.wasm.numThreads = 1;

let instance;
let pending;
let running = false;
const send = message => self.postMessage(message);

/**
 * Interior-photo search. The text tower is bep40/clip-vit-b32-multilingual-vi-text:
 * a CLIP-B/32 text config with mean pooling and the 768->512 projection baked in,
 * so Vietnamese queries land in the same 512-d space as the image matrix
 * (projection_dim: 512 in config.json, confirmed on the Hub).
 *
 * The tokenizer comes from the upstream sentence-transformers repo, NOT the fork:
 * the fork's tokenizer_config.json is a BERT WordPiece config ([CLS]/[SEP],
 * model_max_length 512) that does not match its own CLIP config
 * (max_position_embeddings 77). The fork supplies only the ONNX graph.
 */
function initialize() {
  return instance ??= (async () => {
    send({ status: 'loading', label: 'Đang tải mô hình tìm kiếm' });
    const progress_callback = event => {
      if (event.status === 'progress' && event.file?.endsWith('.onnx')) {
        send({ status: 'loading', label: 'Đang tải mô hình tìm kiếm', progress: event.progress });
      }
    };
    const [tokenizer, model, metadata, buffer] = await Promise.all([
      AutoTokenizer.from_pretrained(TOKENIZER),
      // onnx/ contains exactly one graph, text_model.onnx (130.6 MB), and no
      // *_quantized variant -- verified on the Hub. Any dtype other than 'fp32'
      // would ask for a file that does not exist, so pin 'fp32' and device 'wasm'.
      //
      // `revision` is load-bearing, not a reproducibility nicety: without it the
      // graph is fetched and cached under the unversioned /resolve/main/ URL, so
      // browsers that cached the broken 68dc949 export keep failing with
      // [ShapeInferenceError] Inferred=1 Declared=768 on every reload. See
      // TEXT_MODEL_REVISION in assets.js.
      CLIPTextModelWithProjection.from_pretrained(
        TEXT_MODEL,
        { dtype: 'fp32', device: 'wasm', revision: TEXT_MODEL_REVISION, progress_callback },
      ),
      cachedFile(DATASET + 'image-embeddings.json').then(buffer => JSON.parse(new TextDecoder().decode(buffer))),
      cachedFile(DATASET + 'image-embeddings_25k-512-32bit.bin'),
    ]);
    const embeddings = new Float32Array(buffer);
    if (embeddings.length !== metadata.length * DIMENSIONS) {
      throw new Error('The image database and embeddings do not match.');
    }
    normalizeDatabase(embeddings, DIMENSIONS);
    send({ status: 'loading', label: 'Đang khởi động tìm kiếm', progress: 100 });
    await model(tokenizer('a photograph', { padding: true, truncation: true }));
    send({ status: 'ready', count: metadata.length });
    return { tokenizer, model, metadata, embeddings };
  })().catch(error => { instance = null; throw error; });
}

async function drain() {
  if (running) return;
  running = true;
  try {
    const { tokenizer, model, metadata, embeddings } = await initialize();
    while (pending) {
      const request = pending;
      pending = null;
      const started = performance.now();
      try {
        const { text_embeds } = await model(tokenizer(request.text, { padding: true, truncation: true }));
        // A newer query has already arrived: don't animate an obsolete result.
        if (pending) continue;
        // Keep more interiors than the field shows: merge.js then interleaves
        // the product hits into this list, so a top-60 request that yields 12
        // products still fills a 96-tile field instead of stopping at 60.
        const output = rankImages(text_embeds.data, embeddings, DIMENSIONS, 96)
          .map(({ index, score }) => ({ ...metadata[index], score }));
        send({ status: 'complete', id: request.id, text: request.text, output, elapsed: performance.now() - started });
      } catch (error) {
        send({ status: 'error', id: request.id, message: error.message });
      }
    }
  } catch (error) {
    send({ status: 'error', id: pending?.id, message: error.message });
    pending = null;
  } finally { running = false; }
}

/**
 * Rank Malloca products against the live catalogue that vai-avatar writes.
 *
 * This is deliberately lexical rather than a baked CLIP matrix. A matrix built
 * from a snapshot can only ever know the products that existed when it was
 * built, so "Thêm SP" on vai-avatar would stay invisible until someone re-ran a
 * 6,372-image embed. Reading the same JSON files vai-avatar writes means a new
 * product is searchable on the next query, and a deleted one disappears on its
 * own -- no rebuild, no staleness window.
 *
 * The catalogue is kept warm so a keystroke does not refetch; `refresh-catalogue`
 * forces a re-read, and the periodic poll re-reads it without a query.
 */
let catalogue = null;

function catalogueItems() {
  return catalogue ?? loadCatalogue().then(next => { catalogue = next; return next; });
}

/** Re-read both source files. Shared by the poll loop and the manual refresh. */
async function refreshCatalogue(force = false) {
  const next = force ? await loadCatalogue() : await tryLoadCatalogue();
  if (next) catalogue = next;
  return next;
}

async function searchProducts(text, id, limit) {
  const started = performance.now();
  try {
    send({ status: 'products-loading' });
    const { items, signature } = await catalogueItems();
    const output = searchCatalogue(text, items, limit);
    send({
      status: 'products-complete', id, text, output,
      signature, count: items.size, elapsed: performance.now() - started,
    });
  } catch (error) {
    send({ status: 'products-error', id, message: error.message });
  }
}

self.addEventListener('message', ({ data }) => {
  if (data.type === 'init') {
    initialize().catch(error => send({ status: 'error', message: error.message }));
    // Warm the catalogue in the background so the first product query does not
    // wait on two dataset downloads. A failure here is non-fatal: it surfaces
    // on the first search, where it can be reported properly.
    catalogueItems().catch(() => {});
  } else if (data.type === 'search') {
    pending = { id: data.id, text: data.text };
    drain();
  } else if (data.type === 'search-products') {
    searchProducts(data.text, data.id, data.limit ?? 24);
  } else if (data.type === 'refresh-catalogue') {
    refreshCatalogue(data.force === true)
      .then(next => next && send({ status: 'catalogue-updated', signature: next.signature, count: next.items.size, syncedAt: next.syncedAt }))
      .catch(error => send({ status: 'catalogue-error', message: error.message }));
  } else if (data.type === 'clear') {
    pending = null;
  }
});