import {
  LOCAL_AI_MODELS,
  type LocalAiCapabilities,
  type LocalAiModelDescriptor,
  type LocalAiModelId,
  type LocalAiModelState,
} from '@learnitive/project-ai';

const STORE_DIRECTORY = 'interactive-local-ai-v1';
const MODEL_DIRECTORY = 'models';

type WorkerRequest = {
  id: string;
  type: 'capabilities' | 'model-state' | 'download' | 'remove' | 'cancel';
  modelId?: LocalAiModelId;
};

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

const scope = globalThis as unknown as {
  postMessage(message: WorkerResponse): void;
  onmessage: ((event: MessageEvent<WorkerRequest>) => void) | null;
};
const operationControllers = new Map<string, AbortController>();
const downloadingModelIds = new Set<LocalAiModelId>();

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

function modelDescriptor(modelId: LocalAiModelId | undefined): LocalAiModelDescriptor {
  const model = LOCAL_AI_MODELS.find((item) => item.id === modelId);
  if (!model) throw new Error('Unknown local AI model');
  return model;
}

function storageManager(): StorageManager & { getDirectory?: () => Promise<FileSystemDirectoryHandle> } {
  return navigator.storage as StorageManager & { getDirectory?: () => Promise<FileSystemDirectoryHandle> };
}

async function capabilities(): Promise<LocalAiCapabilities> {
  const webGpu = 'gpu' in navigator && Boolean((navigator as Navigator & { gpu?: unknown }).gpu);
  const opfs = typeof storageManager().getDirectory === 'function';
  return {
    webGpu,
    opfs,
    supported: webGpu && opfs,
    ...(!webGpu
      ? { reason: 'WebGPU is unavailable in this browser.' }
      : !opfs
        ? { reason: 'Origin-private file storage is unavailable in this browser.' }
        : {}),
  };
}

async function modelRoot(create = true): Promise<FileSystemDirectoryHandle> {
  const getDirectory = storageManager().getDirectory;
  if (!getDirectory) throw new Error('Origin-private model storage is unavailable');
  const root = await getDirectory.call(navigator.storage);
  const store = await root.getDirectoryHandle(STORE_DIRECTORY, { create });
  return store.getDirectoryHandle(MODEL_DIRECTORY, { create });
}

function modelFileName(modelId: LocalAiModelId): string {
  return `${modelId}.litertlm`;
}

function metadataFileName(modelId: LocalAiModelId): string {
  return `${modelId}.json`;
}

function opfsWriteChunk(value: Uint8Array<ArrayBufferLike>): ArrayBuffer {
  // FileSystemWritableFileStream.write() is stricter than fetch stream typing
  // under TypeScript's generic typed arrays: OPFS accepts ArrayBuffer-backed
  // data, while ReadableStream chunks are typed as ArrayBufferLike and may in
  // principle be backed by SharedArrayBuffer. Avoid copying the common fetch
  // case where the chunk already owns a full ArrayBuffer; copy only subviews or
  // non-ArrayBuffer-backed chunks.
  const buffer = value.buffer;
  if (buffer instanceof ArrayBuffer) {
    if (value.byteOffset === 0 && value.byteLength === buffer.byteLength) return buffer;
    return buffer.slice(value.byteOffset, value.byteOffset + value.byteLength);
  }
  const copy = new Uint8Array(value.byteLength);
  copy.set(value);
  return copy.buffer;
}

async function removeEntry(directory: FileSystemDirectoryHandle, name: string): Promise<void> {
  await directory.removeEntry(name).catch(() => undefined);
}

async function readMetadata(modelId: LocalAiModelId): Promise<{ size: number } | undefined> {
  try {
    const directory = await modelRoot(false);
    const handle = await directory.getFileHandle(metadataFileName(modelId));
    const file = await handle.getFile();
    const parsed = JSON.parse(await file.text()) as {
      version?: number;
      complete?: boolean;
      modelId?: string;
      modelUrl?: string;
      artifactRevision?: string;
      size?: number;
      installedAt?: number;
    };
    const descriptor = modelDescriptor(modelId);
    if (parsed.complete !== true || parsed.modelId !== modelId || typeof parsed.size !== 'number') return undefined;

    const model = await (await directory.getFileHandle(modelFileName(modelId))).getFile();
    if (model.size !== parsed.size || model.size <= 0) return undefined;

    if (parsed.version === 2) {
      if (parsed.artifactRevision !== descriptor.artifactRevision) return undefined;
      return { size: model.size };
    }

    if (parsed.version !== 1) return undefined;

    // Build Fix 3 migration: v1 installations were bound to an exact URL. A
    // mirror/redirect/descriptor URL change must not throw away a completed
    // multi-GB model. The model id + committed byte size are sufficient to
    // adopt that known-good legacy install into the current artifact revision.
    let metadataWritable: FileSystemWritableFileStream | undefined;
    try {
      metadataWritable = await handle.createWritable({ keepExistingData: false });
      await metadataWritable.write(JSON.stringify({
        version: 2,
        complete: true,
        modelId,
        artifactRevision: descriptor.artifactRevision,
        sourceUrl: descriptor.modelUrl,
        size: model.size,
        installedAt: parsed.installedAt ?? Date.now(),
        migratedFromVersion: 1,
      }));
      await metadataWritable.close();
    } catch {
      // Metadata migration is best-effort. A storage/open/write failure must
      // not turn a previously completed, size-verified v1 model into
      // disposable partial state; keep using it and retry migration later.
      await metadataWritable?.abort().catch(() => undefined);
    }
    return { size: model.size };
  } catch {
    return undefined;
  }
}

async function modelState(modelId: LocalAiModelId): Promise<LocalAiModelState> {
  const metadata = await readMetadata(modelId);
  if (!metadata) {
    // A browser/tab crash can interrupt a FileSystemWritableFileStream before
    // the completion metadata is committed. Treat such bytes as disposable
    // partial state so a failed download never masquerades as an installation
    // or occupies model storage indefinitely.
    const directory = await modelRoot(false).catch(() => undefined);
    if (directory && !downloadingModelIds.has(modelId)) {
      await removeEntry(directory, metadataFileName(modelId));
      await removeEntry(directory, modelFileName(modelId));
    }
    return { modelId, status: 'not-installed' };
  }
  return { modelId, status: 'installed', installedBytes: metadata.size };
}

async function downloadModel(
  requestId: string,
  modelId: LocalAiModelId,
  signal: AbortSignal,
): Promise<LocalAiModelState> {
  const caps = await capabilities();
  if (!caps.opfs) throw new Error(caps.reason ?? 'Local model storage is unavailable');
  if (downloadingModelIds.has(modelId)) throw new Error('This local model is already downloading');
  // Ask the browser to keep the model cache persistent when possible. This is
  // best-effort and, importantly, does not require buffering model bytes in RAM.
  if (!/Firefox\/|FxiOS\//u.test(navigator.userAgent || '')) {
    await navigator.storage.persist?.().catch(() => false);
  }
  downloadingModelIds.add(modelId);
  const model = modelDescriptor(modelId);
  const directory = await modelRoot();
  let writable: FileSystemWritableFileStream | undefined;
  let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
  try {
    await removeEntry(directory, metadataFileName(modelId));
    await removeEntry(directory, modelFileName(modelId));

    const response = await fetch(model.modelUrl, {
      method: 'GET',
      credentials: 'omit',
      redirect: 'follow',
      signal,
    });
    if (!response.ok || !response.body) {
      throw new Error(`Local model download failed (${response.status})`);
    }
    const contentType = response.headers.get('content-type')?.toLowerCase() ?? '';
    if (contentType.includes('text/html') || contentType.includes('application/json')) {
      throw new Error('Local model host returned an unexpected document instead of model bytes');
    }
    const total = Number(response.headers.get('content-length')) || undefined;
    if (total && navigator.storage.estimate) {
      const estimate = await navigator.storage.estimate().catch(() => undefined);
      const available = estimate?.quota !== undefined && estimate.usage !== undefined
        ? Math.max(0, estimate.quota - estimate.usage)
        : undefined;
      if (available !== undefined && available < total + 64 * 1024 * 1024) {
        throw new Error(`Not enough browser storage for this local model. Need about ${Math.ceil(total / (1024 * 1024))} MB plus working space.`);
      }
    }
    const fileHandle = await directory.getFileHandle(modelFileName(modelId), { create: true });
    writable = await fileHandle.createWritable({ keepExistingData: false });
    reader = response.body.getReader();
    let downloaded = 0;
    let lastProgressAt = 0;
    let lastProgressBytes = 0;
    while (true) {
      signal.throwIfAborted();
      const { value, done } = await reader.read();
      if (done) break;
      if (!value?.byteLength) continue;
      // Preserve the zero-copy common case while satisfying OPFS's strict
      // ArrayBuffer-backed write contract under native-preview TypeScript.
      await writable.write(opfsWriteChunk(value));
      downloaded += value.byteLength;
      const now = performance.now();
      if (now - lastProgressAt >= 250 || downloaded - lastProgressBytes >= 4 * 1024 * 1024 || (total && downloaded >= total)) {
        lastProgressAt = now;
        lastProgressBytes = downloaded;
        post({
          id: requestId,
          type: 'progress',
          state: { modelId, status: 'downloading', bytesDownloaded: downloaded, ...(total ? { totalBytes: total } : {}) },
        });
      }
    }
    await writable.close();
    writable = undefined;
    const stored = await fileHandle.getFile();
    if (stored.size <= 0 || (total && stored.size !== total)) {
      throw new Error('Downloaded local model is incomplete');
    }
    const metadataHandle = await directory.getFileHandle(metadataFileName(modelId), { create: true });
    const metadataWritable = await metadataHandle.createWritable({ keepExistingData: false });
    try {
      await metadataWritable.write(JSON.stringify({
        version: 2,
        complete: true,
        modelId,
        artifactRevision: model.artifactRevision,
        sourceUrl: model.modelUrl,
        size: stored.size,
        installedAt: Date.now(),
      }));
      await metadataWritable.close();
    } catch (error) {
      await metadataWritable.abort().catch(() => undefined);
      throw error;
    }
    return { modelId, status: 'installed', installedBytes: stored.size };
  } catch (error) {
    await writable?.abort().catch(() => undefined);
    await removeEntry(directory, metadataFileName(modelId));
    await removeEntry(directory, modelFileName(modelId));
    throw error;
  } finally {
    reader?.releaseLock();
    downloadingModelIds.delete(modelId);
  }
}

scope.onmessage = (event) => {
  const request = event.data;
  if (request.type === 'cancel') {
    operationControllers.get(request.id)?.abort();
    return;
  }
  const controller = new AbortController();
  operationControllers.set(request.id, controller);
  void (async () => {
    try {
      let value: unknown;
      if (request.type === 'capabilities') value = await capabilities();
      else if (request.type === 'model-state') value = await modelState(modelDescriptor(request.modelId).id);
      else if (request.type === 'download') value = await downloadModel(request.id, modelDescriptor(request.modelId).id, controller.signal);
      else if (request.type === 'remove') {
        const modelId = modelDescriptor(request.modelId).id;
        const directory = await modelRoot().catch(() => undefined);
        if (directory) {
          await removeEntry(directory, modelFileName(modelId));
          await removeEntry(directory, metadataFileName(modelId));
        }
      }
      post({ id: request.id, type: 'response', ok: true, value });
    } catch (error) {
      const message = controller.signal.aborted
        ? 'The local AI operation was cancelled.'
        : error instanceof Error
          ? error.message
          : String(error);
      post({ id: request.id, type: 'response', ok: false, error: message });
    } finally {
      operationControllers.delete(request.id);
    }
  })();
};