import {
  EMBEDDING_DIMENSIONS,
  EMBEDDING_MODEL_ID,
  type EmbeddingCapabilities,
  type EmbeddingModelState,
} from '@learnitive/project-meaning';

const TRANSFORMERS_MODULE_URL = 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1';
const MODEL_ID = 'onnx-community/embeddinggemma-300m-ONNX';
const MODEL_REVISION = '75a84c732f1884df76bec365346230e32f582c82';
const MODEL_VERSION = 'q8';
const META_DB = 'interactive-embedding-model-v1';
const META_STORE = 'models';

type WorkerRequest = {
  id: string;
  type: 'capabilities' | 'model-state' | 'prepare-model' | 'remove-model' | 'load' | 'unload' | 'embed-documents' | 'embed-query' | 'cancel';
  texts?: readonly string[];
  query?: string;
  allowDownload?: boolean;
};

type WorkerResponse =
  | { id: string; type: 'response'; ok: true; value?: unknown }
  | { id: string; type: 'response'; ok: false; error: string }
  | { id: string; type: 'progress'; state: EmbeddingModelState };

type Tensor = { data: Float32Array | number[]; dims?: readonly number[] };
type Tokenizer = (texts: string | readonly string[], options?: Record<string, unknown>) => Promise<Record<string, unknown>> | Record<string, unknown>;
type Model = ((inputs: Record<string, unknown>) => Promise<{ sentence_embedding?: Tensor }>) & { dispose?: () => Promise<void> | void };
type TransformersModule = {
  AutoTokenizer: { from_pretrained(modelId: string, options?: Record<string, unknown>): Promise<Tokenizer> };
  AutoModel: { from_pretrained(modelId: string, options?: Record<string, unknown>): Promise<Model> };
  env?: {
    backends?: {
      onnx?: {
        wasm?: {
          wasmPaths?: string;
          numThreads?: number;
        };
      };
    };
  };
};

const scope = globalThis as unknown as {
  postMessage(message: WorkerResponse): void;
  onmessage: ((event: MessageEvent<WorkerRequest>) => void) | null;
};
const controllers = new Map<string, AbortController>();
let modulePromise: Promise<TransformersModule> | undefined;
let tokenizer: Tokenizer | undefined;
let model: Model | undefined;
let loading: Promise<void> | undefined;

function post(message: WorkerResponse): void {
  scope.postMessage(message);
}

function openMetaDb(): Promise<IDBDatabase> {
  return new Promise((resolve, reject) => {
    const request = indexedDB.open(META_DB, 1);
    request.onupgradeneeded = () => {
      if (!request.result.objectStoreNames.contains(META_STORE)) request.result.createObjectStore(META_STORE);
    };
    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error ?? new Error('Embedding metadata database failed'));
  });
}

async function readInstalled(): Promise<boolean> {
  const db = await openMetaDb();
  try {
    return await new Promise<boolean>((resolve, reject) => {
      const tx = db.transaction(META_STORE, 'readonly');
      const request = tx.objectStore(META_STORE).get(EMBEDDING_MODEL_ID);
      request.onsuccess = () => {
        const value = request.result as { modelId?: string; version?: string; installed?: boolean } | undefined;
        resolve(value?.installed === true && value.modelId === MODEL_ID && value.version === MODEL_VERSION);
      };
      request.onerror = () => reject(request.error ?? new Error('Embedding metadata read failed'));
    });
  } finally {
    db.close();
  }
}

async function writeInstalled(installed: boolean): Promise<void> {
  const db = await openMetaDb();
  try {
    await new Promise<void>((resolve, reject) => {
      const tx = db.transaction(META_STORE, 'readwrite');
      const store = tx.objectStore(META_STORE);
      const request = installed
        ? store.put({ modelId: MODEL_ID, version: MODEL_VERSION, installed: true, installedAt: Date.now() }, EMBEDDING_MODEL_ID)
        : store.delete(EMBEDDING_MODEL_ID);
      request.onsuccess = () => resolve();
      request.onerror = () => reject(request.error ?? new Error('Embedding metadata write failed'));
    });
  } finally {
    db.close();
  }
}

async function capabilities(): Promise<EmbeddingCapabilities> {
  const webGpu = 'gpu' in navigator && Boolean((navigator as Navigator & { gpu?: unknown }).gpu);
  const wasm = typeof WebAssembly !== 'undefined';
  return {
    supported: wasm,
    webGpu,
    wasm,
    ...(!wasm ? { reason: 'WebAssembly is unavailable in this browser.' } : {}),
  };
}

async function transformers(): Promise<TransformersModule> {
  modulePromise ??= (import(/* @vite-ignore */ TRANSFORMERS_MODULE_URL) as Promise<TransformersModule>)
    .then((module) => {
      // Transformers.js defaults to a CDN for ONNX Runtime WASM assets. Be
      // explicit here so a Vite/module-worker URL cannot be mistaken for the
      // ORT asset root. A single WASM thread also avoids spawning another
      // cross-origin worker merely to embed text.
      const wasm = module.env?.backends?.onnx?.wasm;
      if (wasm) {
        wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.22.0-dev.20250409-89f8206ba4/dist/';
        wasm.numThreads = 1;
      }
      return module;
    })
    .catch((error) => {
      modulePromise = undefined;
      throw error;
    });
  return modulePromise;
}

function progressState(progress: unknown): EmbeddingModelState | undefined {
  if (!progress || typeof progress !== 'object') return undefined;
  const value = progress as { progress?: number; loaded?: number; total?: number; status?: string };
  const fraction = typeof value.progress === 'number'
    ? (value.progress > 1 ? value.progress / 100 : value.progress)
    : undefined;
  return {
    modelId: EMBEDDING_MODEL_ID,
    status: 'downloading',
    ...(fraction !== undefined ? { progress: Math.max(0, Math.min(1, fraction)) } : {}),
    ...(typeof value.loaded === 'number' ? { loadedBytes: value.loaded } : {}),
    ...(typeof value.total === 'number' ? { totalBytes: value.total } : {}),
  };
}

async function deleteCachedModelAssets(): Promise<void> {
  if (typeof caches === 'undefined') return;
  const needles = [MODEL_ID, encodeURIComponent(MODEL_ID), 'embeddinggemma-300m-ONNX'];
  for (const cacheName of await caches.keys()) {
    const cache = await caches.open(cacheName);
    for (const request of await cache.keys()) {
      if (needles.some((needle) => request.url.includes(needle))) await cache.delete(request);
    }
  }
}

async function loadModel(requestId: string, allowDownload: boolean, signal: AbortSignal): Promise<void> {
  if (model && tokenizer) return;
  if (loading) return loading;
  const installed = await readInstalled();
  if (!allowDownload && !installed) throw new Error('Download EmbeddingGemma before using Meaning search');
  const caps = await capabilities();
  if (!caps.supported) throw new Error(caps.reason ?? 'Local embeddings are unavailable');

  loading = (async () => {
    signal.throwIfAborted();
    const module = await transformers();
    const common = {
      local_files_only: !allowDownload,
      revision: MODEL_REVISION,
      progress_callback: (value: unknown) => {
        if (signal.aborted) throw signal.reason ?? new DOMException('The operation was aborted', 'AbortError');
        if (!allowDownload) return;
        const state = progressState(value);
        if (state) post({ id: requestId, type: 'progress', state });
      },
    };
    try {
      const nextTokenizer = await module.AutoTokenizer.from_pretrained(MODEL_ID, common);
      signal.throwIfAborted();

      // Embeddings do not need the large-model WebGPU path. Prefer the
      // broadly-supported ONNX/WASM backend and fall back to WebGPU only when
      // WASM genuinely cannot initialize. This also avoids the confusing
      // `no available backend found` failure seen after a successful download
      // on browsers whose WebGPU ONNX backend cannot load its JSEP module.
      const devices = caps.webGpu ? (['wasm', 'webgpu'] as const) : (['wasm'] as const);
      let nextModel: Model | undefined;
      const backendErrors: string[] = [];
      for (const device of devices) {
        try {
          nextModel = await module.AutoModel.from_pretrained(MODEL_ID, {
            ...common,
            dtype: 'q8',
            device,
          });
          break;
        } catch (error) {
          backendErrors.push(`${device}: ${error instanceof Error ? error.message : String(error)}`);
        }
      }
      if (!nextModel) {
        throw new Error(`EmbeddingGemma could not initialize a browser inference backend. ${backendErrors.join(' | ')}`);
      }
      signal.throwIfAborted();
      tokenizer = nextTokenizer;
      model = nextModel;
      if (allowDownload) await writeInstalled(true);
    } catch (error) {
      tokenizer = undefined;
      model = undefined;
      if (!allowDownload && installed) await writeInstalled(false).catch(() => undefined);
      if (allowDownload) {
        // Do not purge hundreds of megabytes of successfully downloaded model
        // shards merely because a browser backend failed to initialize. The
        // Cache API remains reusable for the WASM/WebGPU retry or next load.
        await writeInstalled(false).catch(() => undefined);
      }
      throw error;
    }
  })().finally(() => {
    loading = undefined;
  });
  return loading;
}

async function unload(): Promise<void> {
  await model?.dispose?.();
  model = undefined;
  tokenizer = undefined;
}

function normalizeRows(tensor: Tensor, expectedRows: number): Float32Array[] {
  const source = tensor.data instanceof Float32Array ? tensor.data : Float32Array.from(tensor.data);
  const dimensions = tensor.dims?.at(-1) ?? EMBEDDING_DIMENSIONS;
  if (dimensions <= 0 || source.length !== expectedRows * dimensions) {
    throw new Error(`Unexpected EmbeddingGemma output shape (${source.length} values for ${expectedRows} inputs)`);
  }
  const rows: Float32Array[] = [];
  for (let row = 0; row < expectedRows; row += 1) {
    const vector = source.slice(row * dimensions, (row + 1) * dimensions);
    let norm = 0;
    for (const value of vector) norm += value * value;
    norm = Math.sqrt(norm);
    if (norm > 0) for (let i = 0; i < vector.length; i += 1) vector[i] = vector[i]! / norm;
    rows.push(vector);
  }
  return rows;
}

async function embed(texts: readonly string[], kind: 'query' | 'document', signal: AbortSignal): Promise<Float32Array[]> {
  if (!model || !tokenizer) throw new Error('EmbeddingGemma is not loaded');
  signal.throwIfAborted();
  const prefix = kind === 'query' ? 'task: search result | query: ' : 'title: none | text: ';
  const prepared = texts.map((text) => `${prefix}${text}`);
  const inputs = await tokenizer(prepared, { padding: true, truncation: true, max_length: 2048 });
  signal.throwIfAborted();
  const result = await model(inputs);
  signal.throwIfAborted();
  if (!result.sentence_embedding) throw new Error('EmbeddingGemma did not return sentence embeddings');
  return normalizeRows(result.sentence_embedding, prepared.length);
}

async function purgeModelCache(): Promise<void> {
  await unload();
  await writeInstalled(false);
  await deleteCachedModelAssets();
}

scope.onmessage = (event) => {
  const request = event.data;
  if (request.type === 'cancel') {
    controllers.get(request.id)?.abort();
    return;
  }
  const controller = new AbortController();
  controllers.set(request.id, controller);
  void (async () => {
    try {
      let value: unknown;
      if (request.type === 'capabilities') value = await capabilities();
      else if (request.type === 'model-state') {
        const installed = await readInstalled();
        value = { modelId: EMBEDDING_MODEL_ID, status: model && tokenizer ? 'ready' : installed ? 'installed' : 'not-installed' } satisfies EmbeddingModelState;
      } else if (request.type === 'prepare-model') {
        post({ id: request.id, type: 'progress', state: { modelId: EMBEDDING_MODEL_ID, status: 'loading' } });
        await loadModel(request.id, true, controller.signal);
        value = { modelId: EMBEDDING_MODEL_ID, status: 'ready', progress: 1 } satisfies EmbeddingModelState;
      } else if (request.type === 'remove-model') {
        await purgeModelCache();
      } else if (request.type === 'load') {
        await loadModel(request.id, request.allowDownload === true, controller.signal);
      } else if (request.type === 'unload') {
        await unload();
      } else if (request.type === 'embed-documents') {
        const texts = request.texts ?? [];
        if (texts.length === 0) value = [];
        else {
          await loadModel(request.id, false, controller.signal);
          value = await embed(texts, 'document', controller.signal);
        }
      } else if (request.type === 'embed-query') {
        const query = request.query?.trim();
        if (!query) throw new Error('Meaning search query is required');
        await loadModel(request.id, false, controller.signal);
        value = (await embed([query], 'query', controller.signal))[0];
      }
      post({ id: request.id, type: 'response', ok: true, value });
    } catch (error) {
      post({ id: request.id, type: 'response', ok: false, error: error instanceof Error ? error.message : String(error) });
    } finally {
      controllers.delete(request.id);
    }
  })();
};