/**
 * Inference Web Worker (Decisions.md D-003 — inference never runs on the main
 * thread). Hosts the Transformers.js `automatic-speech-recognition` pipeline.
 *
 * M2: hardware backend detection with automatic fallback — "auto" resolves
 * WebNN → WebGPU → WASM (D-004, D-017) — catalog-driven dtype defaults
 * (D-006), model switching with pipeline disposal, and language/task options
 * for multilingual models.
 */

import {
  env,
  pipeline,
  WhisperTextStreamer,
  type AutomaticSpeechRecognitionPipeline,
  type ProgressInfo,
} from "@huggingface/transformers";
import { resolveDefaultDtype } from "../lib/models";
import {
  deviceAttempts,
  dtypeRequiresF16,
  isUiToWorkerMessage,
  TranscribeProgressTracker,
  type DeviceCapabilities,
  type DevicePreference,
  type InferenceDevice,
  type LoadMessage,
  type ResultMessage,
  type TranscribeMessage,
  type TranscriptChunk,
  type WorkerToUiMessage,
} from "./protocol";

// Load onnxruntime's .wasm binaries from our own origin instead of the
// jsdelivr CDN default — required for offline use (M3). The server exposes
// them at /ort/* (src/index.ts; copied into dist/ort by scripts/build.ts)
// and the service worker runtime-caches them cache-first (src/pwa/sw.ts).
if (env.backends.onnx.wasm) {
  env.backends.onnx.wasm.wasmPaths = "./ort/";
}

/**
 * Set the WASM thread count for the device we are about to create a pipeline
 * on (D-023). The server sends COOP/COEP, so SharedArrayBuffer exists and
 * threads genuinely run — which cuts both ways:
 * - Pure WASM inference scales well past onnxruntime-web's 4-thread default,
 *   so use most of the machine (measured 52 s → 9.5 s on a 14 s clip).
 * - The JSEP-based accelerator EPs (WebGPU, WebNN) are NOT thread-safe:
 *   with numThreads > 1 (even ort's own default of ~4 once SAB exists),
 *   inference fails with "WebGPU validation failed. Buffer with '' label is
 *   invalid". They must run the wasm side single-threaded.
 * Caveat: onnxruntime-web initializes its runtime on the first session
 * creation and ignores later numThreads changes, so the first attempt in
 * this worker decides. A wasm *fallback* after a failed accelerator attempt
 * therefore may run single-threaded — correct, just slower; the next worker
 * respawn (cancel does one, D-018) picks the full count up again.
 */
function desiredWasmThreads(): number {
  return globalThis.crossOriginIsolated
    ? Math.max(1, Math.min((navigator.hardwareConcurrency || 4) - 1, 12))
    : 1;
}

function configureWasmThreads(device: InferenceDevice): void {
  const wasm = env.backends.onnx.wasm;
  if (!wasm) return;
  wasm.numThreads = device === "wasm" ? desiredWasmThreads() : 1;
}

/**
 * Thrown to abort a load whose WASM fallback would be stuck single-threaded
 * (D-025). The worker has already posted `load-restart`; the UI terminates
 * this worker and retries in a fresh one, so nothing is reported as an error.
 */
class RestartRequested extends Error {
  constructor() {
    super("load aborted pending worker restart (D-025)");
  }
}

// `lib: ["DOM"]` types `self` as Window; give it the worker-scope shape we use.
interface DedicatedWorkerScope {
  postMessage(message: WorkerToUiMessage): void;
  addEventListener(
    type: "message",
    listener: (event: MessageEvent<unknown>) => void,
  ): void;
}
const scope = self as unknown as DedicatedWorkerScope;

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

function errorMessage(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}

/**
 * Forward Transformers.js progress events to the UI as load-progress messages.
 * The library's own "ready" event is swallowed — the worker posts its own
 * `ready` (with the selected device) once the pipeline actually resolved,
 * which also covers the device fallback retries (D-004, D-017).
 */
function forwardProgress(info: ProgressInfo): void {
  switch (info.status) {
    case "initiate":
    case "download":
    case "done":
      post({ type: "load-progress", status: info.status, file: info.file });
      break;
    case "progress":
      post({
        type: "load-progress",
        status: "progress",
        file: info.file,
        loaded: info.loaded,
        total: info.total,
        progress: info.progress,
      });
      break;
    default:
      // "ready" (posted by the worker itself, with the device) and
      // "progress_total" — the UI derives totals per-file itself.
      break;
  }
}

// ---------------------------------------------------------------------------
// Device selection (D-004, D-017): "auto" resolves WebNN → WebGPU → WASM.
// The probe only answers "does the backend exist?"; the attempt ordering and
// fallback policy live in the pure `deviceAttempts` (protocol.ts, unit-tested).
// ---------------------------------------------------------------------------

interface GpuAdapterLike {
  features?: { has(name: string): boolean };
  /** Pre-2024 spec location of the software-adapter flag. */
  isFallbackAdapter?: boolean;
  /** Current spec location (GPUAdapterInfo). */
  info?: { isFallbackAdapter?: boolean };
}

/** True for software adapters (SwiftShader, llvmpipe, WARP) — GPU-shaped but
 * slower than multithreaded WASM, so never worth an attempt (D-025). */
function isSoftwareAdapter(adapter: GpuAdapterLike): boolean {
  return Boolean(adapter.isFallbackAdapter ?? adapter.info?.isFallbackAdapter);
}

interface NavigatorWithHardware {
  gpu?: { requestAdapter(): Promise<GpuAdapterLike | null> };
  ml?: unknown;
}

async function probeCapabilities(
  preference: DevicePreference,
): Promise<DeviceCapabilities> {
  const capabilities: DeviceCapabilities = { webnn: false, webgpu: false };
  const hardware = navigator as NavigatorWithHardware;
  if (preference === "auto" || preference === "webnn") {
    try {
      capabilities.webnn = Boolean(hardware.ml);
    } catch {
      // Probing threw — treat WebNN as unavailable.
    }
  }
  if (preference === "auto" && /Firefox\//.test(navigator.userAgent)) {
    // Firefox creates ort-web WebGPU sessions fine but runs them far slower
    // than this machine's multithreaded WASM (field-verified 2026-07-17,
    // D-026) — and a *successful* slow session is invisible to the
    // failure-driven fallback chain. Never auto-pick WebGPU on Firefox; an
    // explicit selection is honored (the retest hatch for future ort/Firefox).
    console.info(
      "[whisper] Firefox: skipping WebGPU for \"auto\" — onnxruntime-web's WebGPU backend " +
        "runs much slower than multithreaded WASM on Firefox (D-026). " +
        "Select WebGPU explicitly to retest.",
    );
  } else if (preference === "auto" || preference === "webgpu") {
    try {
      const adapter = (await hardware.gpu?.requestAdapter()) ?? null;
      if (adapter && isSoftwareAdapter(adapter)) {
        // A fallback (software-rendered) adapter would "succeed" and then run
        // far slower than multithreaded WASM — never attempt it (D-025).
        console.info("[whisper] WebGPU adapter is software-rendered — skipping WebGPU.");
      } else {
        capabilities.webgpu = adapter !== null;
        // fp16 weights need the "shader-f16" feature (D-025); Firefox ships
        // WebGPU without it, so createPipeline gates on this.
        capabilities.webgpuF16 = Boolean(adapter?.features?.has("shader-f16"));
      }
    } catch {
      // Probing threw — treat WebGPU as unavailable.
    }
  }
  return capabilities;
}

// ---------------------------------------------------------------------------
// Pipeline lifecycle (load, switch, dispose)
// ---------------------------------------------------------------------------

interface ActivePipeline {
  /** Identity of this load request; a different key means "switch models". */
  key: string;
  promise: Promise<AutomaticSpeechRecognitionPipeline>;
}

let active: ActivePipeline | null = null;

function loadKey(message: LoadMessage): string {
  return JSON.stringify({
    modelId: message.modelId,
    device: message.device,
    dtype: message.dtype ?? null,
  });
}

async function createPipeline(
  message: LoadMessage,
): Promise<AutomaticSpeechRecognitionPipeline> {
  const capabilities = await probeCapabilities(message.device);
  // WebGPU without "shader-f16" (Firefox) cannot run fp16 weights — skip the
  // doomed attempt instead of failing into a degraded-thread fallback (D-025).
  if (capabilities.webgpu && !capabilities.webgpuF16) {
    const webgpuDtype =
      message.dtype ?? resolveDefaultDtype(message.modelId, "webgpu");
    if (dtypeRequiresF16(webgpuDtype)) {
      capabilities.webgpu = false;
      console.info(
        "[whisper] WebGPU adapter lacks shader-f16 but the model dtype needs fp16 — skipping WebGPU.",
      );
    }
  }
  // Ordered devices to try (D-004/D-017): the first is the selected device;
  // if pipeline creation throws there (driver quirks, unsupported ops,
  // out-of-memory), the next one is tried automatically. For "auto" the chain
  // is WebNN → WebGPU → WASM (filtered by the probe); WASM always terminates it.
  const attempts: InferenceDevice[] = deviceAttempts(message.device, capabilities);

  let lastError: unknown;
  let acceleratorAttempted = false;
  for (const device of attempts) {
    // A failed WebGPU/WebNN attempt has already initialized onnxruntime's
    // WASM runtime single-threaded, and the thread pool is fixed at first
    // init (D-023) — a WASM session created now would run on one thread for
    // this worker's whole life (measured 5.5× slower). When multithreading is
    // actually available, hand the fallback to a fresh worker instead (D-025).
    if (device === "wasm" && acceleratorAttempted && desiredWasmThreads() > 1) {
      post({
        type: "load-restart",
        modelId: message.modelId,
        preference: message.device,
        device: "wasm",
      });
      throw new RestartRequested();
    }
    configureWasmThreads(device);
    const dtype = message.dtype ?? resolveDefaultDtype(message.modelId, device);
    try {
      const asr = await pipeline("automatic-speech-recognition", message.modelId, {
        device,
        dtype,
        // WASM sessions cap graph optimization at "basic" (D-014): the level-2
        // QDQ transform in the bundled onnxruntime-web build rejects the 8-bit
        // decoders' tied embed_tokens weight ("Missing required scale …
        // DequantizeLinear"). q4 models use native MatMulNBits ops, so the
        // lower level costs them nothing. WebGPU keeps the default level.
        ...(device === "wasm" && {
          session_options: { graphOptimizationLevel: "basic" as const },
        }),
        progress_callback: forwardProgress,
      });
      // What ort actually ended up with (it clamps to 1 without isolation) —
      // surfaced to the UI so degraded threading is visible in the field.
      const threads = env.backends.onnx.wasm?.numThreads ?? 1;
      console.info(`[whisper] pipeline ready: device=${device} wasmThreads=${threads}`);
      post({ type: "load-progress", status: "ready", device, threads });
      return asr;
    } catch (error) {
      lastError = error;
      if (device !== "wasm") acceleratorAttempted = true;
      console.warn(
        `[whisper] pipeline creation failed on ${device}: ${errorMessage(error)}`,
      );
    }
  }
  throw lastError;
}

function load(message: LoadMessage): Promise<AutomaticSpeechRecognitionPipeline> {
  const key = loadKey(message);
  // Idempotent: repeated "load" messages for the same model/device/dtype reuse
  // the in-flight or completed pipeline.
  if (active?.key === key) return active.promise;

  const previous = active;
  const promise = (async () => {
    if (previous) {
      // Dispose the old pipeline (frees WASM memory / GPU buffers) before
      // creating the new one, so both never coexist.
      try {
        const old = await previous.promise;
        await old.dispose();
      } catch {
        // The previous load itself failed — nothing to dispose.
      }
    }
    return createPipeline(message);
  })();
  active = { key, promise };
  promise.catch(() => {
    // Allow a clean retry after a failed load; ignore stale failures that
    // were already superseded by a newer load request.
    if (active?.key === key) active = null;
  });
  return promise;
}

// ---------------------------------------------------------------------------
// Transcription
// ---------------------------------------------------------------------------

/** Audio sample rate the whole app standardizes on (D-012). */
const SAMPLE_RATE = 16_000;
/** Chunked-pipeline window / overlap, in seconds (Whisper's usual 30/5). */
const CHUNK_LENGTH_S = 30;
const STRIDE_LENGTH_S = 5;

async function transcribe(message: TranscribeMessage): Promise<void> {
  if (active === null) {
    post({
      type: "error",
      stage: "transcribe",
      message: "No model loaded — send a \"load\" message first.",
      requestId: message.requestId,
    });
    return;
  }
  const asr = await active.promise;

  const options: Record<string, unknown> = {
    chunk_length_s: CHUNK_LENGTH_S,
    stride_length_s: STRIDE_LENGTH_S,
    return_timestamps: true,
    // With return_timestamps and no explicit max_new_tokens, Transformers.js
    // 4.2's Whisper generate() takes a seek-loop path that re-decodes each
    // window's tail (1–3 generate passes per window). Under this pipeline's
    // own 30 s/5 s chunking that is pure overhead: ~1.4× slower on long
    // files, duplicated boundary text, on_finalize firing once per pass
    // (breaking TranscribeProgressTracker's one-finalize-per-window
    // assumption — the progress bar saturated early, then looked frozen),
    // and a potential infinite loop when a window's last boundary timestamp
    // is 0. Setting max_new_tokens selects the classic single-pass path;
    // 440 ≈ Whisper's 448 decoder positions minus prompt headroom, far more
    // than any real 30 s window produces.
    max_new_tokens: 440,
  };
  if (message.stream ?? message.requestId === undefined) {
    // File-flow pass (id-less, or routed through BrowserEngine with an
    // explicit stream flag, D-010): stream partial tokens and report audio
    // progress. The streamer's chunk timestamps are window-relative; the
    // tracker maps them to monotone absolute seconds (see protocol.ts).
    // Live passes skip both — they resolve wholesale via "result" (D-016).
    const totalSeconds = message.audio.length / SAMPLE_RATE;
    const tracker = new TranscribeProgressTracker(
      totalSeconds,
      CHUNK_LENGTH_S,
      CHUNK_LENGTH_S - 2 * STRIDE_LENGTH_S,
    );
    const postProgress = (processedSeconds: number | null) => {
      if (processedSeconds === null) return;
      post({
        type: "transcribe-progress",
        processedSeconds,
        totalSeconds,
        requestId: message.requestId,
      });
    };
    // Immediate 0/total so the UI shows the audio length before the first
    // (possibly slow) token arrives.
    post({
      type: "transcribe-progress",
      processedSeconds: 0,
      totalSeconds,
      requestId: message.requestId,
    });
    options.streamer = new WhisperTextStreamer(
      // The ASR pipeline for Whisper models always carries a WhisperTokenizer.
      asr.tokenizer as ConstructorParameters<typeof WhisperTextStreamer>[0],
      {
        skip_prompt: true,
        callback_function: (text: string) =>
          post({ type: "partial", text, requestId: message.requestId }),
        on_chunk_start: (time: number) => postProgress(tracker.timestamp(time)),
        on_chunk_end: (time: number) => postProgress(tracker.timestamp(time)),
        on_finalize: () => postProgress(tracker.finalizeWindow()),
      },
    );
  }
  // Only meaningful for multilingual models; the UI omits these for `.en`
  // variants. "auto" / undefined lets Whisper detect the language.
  if (message.language !== undefined && message.language !== "auto") {
    options.language = message.language;
  }
  if (message.task !== undefined) {
    options.task = message.task;
  }

  const output = await asr(message.audio, options);

  const single = (Array.isArray(output) ? output[0] : output) as {
    text?: string;
    chunks?: TranscriptChunk[];
  };
  const result: ResultMessage = {
    type: "result",
    text: (single.text ?? "").trim(),
    chunks: single.chunks ?? [],
    requestId: message.requestId,
  };
  post(result);
}

scope.addEventListener("message", (event) => {
  const data = event.data;
  if (!isUiToWorkerMessage(data)) {
    post({
      type: "error",
      stage: "transcribe",
      message: "Malformed message sent to inference worker.",
    });
    return;
  }
  switch (data.type) {
    case "load":
      load(data).catch((error: unknown) => {
        // A restart-aborted load is not an error — the UI already got the
        // load-restart message and is respawning this worker (D-025).
        if (error instanceof RestartRequested) return;
        post({ type: "error", stage: "load", message: errorMessage(error) });
      });
      break;
    case "transcribe":
      transcribe(data).catch((error: unknown) => {
        if (error instanceof RestartRequested) return;
        post({
          type: "error",
          stage: "transcribe",
          message: errorMessage(error),
          requestId: data.requestId,
        });
      });
      break;
  }
});