/**
 * 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,
  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/";
}

// `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 NavigatorWithHardware {
  gpu?: { requestAdapter(): Promise<unknown | 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" || preference === "webgpu") {
    try {
      capabilities.webgpu = Boolean(await hardware.gpu?.requestAdapter());
    } 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> {
  // 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,
    await probeCapabilities(message.device),
  );

  let lastError: unknown;
  for (const device of attempts) {
    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,
      });
      post({ type: "load-progress", status: "ready", device });
      return asr;
    } catch (error) {
      lastError = 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,
  };
  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) => {
        post({ type: "error", stage: "load", message: errorMessage(error) });
      });
      break;
    case "transcribe":
      transcribe(data).catch((error: unknown) => {
        post({
          type: "error",
          stage: "transcribe",
          message: errorMessage(error),
          requestId: data.requestId,
        });
      });
      break;
  }
});