import { AutoTokenizer, CLIPTextModelWithProjection, env } from '@huggingface/transformers';
import { cachedFile, DATASET, TOKENIZER, TEXT_MODEL, DIMENSIONS } from './assets.js';
import { normalizeDatabase, rankImages } from './search.js';
import { PRODUCT_EMBEDDINGS, PRODUCT_ROWS, loadGallery } from './products.js';

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

let instance;
let productsPromise;
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'.
      CLIPTextModelWithProjection.from_pretrained(TEXT_MODEL, { dtype: 'fp32', device: 'wasm', 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; });
}

/** Expand IEEE 754 binary16 into binary32; JS has no widely-supported Float16Array. */
function halfToFloat(view, count) {
  const out = new Float32Array(count);
  for (let i = 0; i < count; i++) {
    const bits = view.getUint16(i * 2, true);
    const sign = (bits & 0x8000) ? -1 : 1;
    const exponent = (bits >> 10) & 0x1f;
    const fraction = bits & 0x03ff;
    let value;
    if (exponent === 0) value = fraction * 2 ** -24;                 // subnormal
    else if (exponent === 0x1f) value = fraction ? NaN : Infinity;   // inf / NaN
    else value = (1 + fraction / 1024) * 2 ** (exponent - 15);
    out[i] = sign * value;
  }
  return out;
}

/**
 * Malloca product photos, ranked in the same text tower -- so a Vietnamese query
 * works directly, with no translation model. The matrix is L2-normalized offline,
 * so ranking is a plain dot product.
 *
 * `meta.rows` is [{ sku, image, index }] written by tools/build_product_embeddings.py,
 * with `index` the row offset in the matrix. It is the single source of truth for
 * the image -> product mapping; nothing here guesses gallery structure.
 */
function initializeProducts() {
  return productsPromise ??= (async () => {
    const on_progress = (loaded, total) => {
      send({ status: 'products-loading', progress: total ? Math.round((loaded / total) * 100) : 0 });
    };
    const [meta, buffer] = await Promise.all([
      fetch(PRODUCT_ROWS).then(response => {
        if (!response.ok) throw new Error(`Không tải được product-emb.json (${response.status}).`);
        return response.json();
      }),
      cachedFile(PRODUCT_EMBEDDINGS, on_progress).catch(error => {
        throw new Error(`Chưa có ma trận embedding sản phẩm. Chạy tools/run_embedding_job.py trước. (${error.message})`);
      }),
    ]);

    const { count, dim, dtype } = meta;
    if (dim !== DIMENSIONS) throw new Error(`Sai kích thước vector: ${dim} != ${DIMENSIONS}`);
    if (!Array.isArray(meta.rows) || meta.rows.length !== count) {
      throw new Error(`product-emb.json: ${meta.rows?.length ?? 0} rows cho ${count} embeddings.`);
    }
    const bytes = dtype === 'float16' ? 2 : 4;
    const expect = count * dim * bytes;
    if (buffer.byteLength !== expect) {
      throw new Error(`Ma trận embedding không khớp: ${buffer.byteLength} B, cần ${expect} B.`);
    }
    const data = dtype === 'float16'
      ? halfToFloat(new DataView(buffer), count * dim)
      : new Float32Array(buffer.slice(0));
    normalizeDatabase(data, dim);
    return { data, dim, count, rows: meta.rows };
  })().catch(error => { productsPromise = 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;
        const output = rankImages(text_embeds.data, embeddings, DIMENSIONS)
          .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; }
}

async function embedText(text) {
  const { tokenizer, model } = await initialize();
  const { text_embeds } = await model(tokenizer(text, { padding: true, truncation: true }));
  const query = text_embeds.data;
  let norm = 0;
  for (const value of query) norm += value * value;
  return { query, norm: Math.sqrt(norm) || 1 };
}

/** Rank Malloca products. Independent of the interior index. */
async function searchProducts(text, id, limit) {
  const started = performance.now();
  try {
    const { data, dim, count, rows } = await initializeProducts();
    const { query, norm } = await embedText(text);

    // A product may have several photos; keep its best-scoring row.
    const best = new Map();
    for (const { index, sku, image } of rows) {
      const offset = index * dim;
      let dot = 0;
      for (let j = 0; j < dim; j++) dot += query[j] * data[offset + j];
      const score = dot / norm;
      const seen = best.get(sku);
      if (!seen || score > seen.score) best.set(sku, { score, image });
    }
    const output = Array.from(best, ([sku, { score, image }]) => ({ sku, score, image }))
      .sort((a, b) => b.score - a.score)
      .slice(0, limit);
    send({ status: 'products-complete', id, text, output, count, 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 }));
  } 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 === 'clear') {
    pending = null;
  }
});