import { describe, test, expect } from "bun:test";
import {
  DEFAULT_MODEL_ID,
  MODELS,
  WHISPER_LANGUAGES,
  estimateDefaultDownloadBytes,
  estimateDownloadBytes,
  estimatePeakLoadBytes,
  formatModelSize,
  getFasterWhisperModel,
  getModel,
  resolveDefaultDtype,
  taskForModel,
} from "../src/lib/models";

// Pure catalog data (D-006) — no browser, no network.

/** Hub repo-name decorations that follow the Whisper size suffix. */
function stripRepoDecorations(id: string): string {
  return id.replace(/(_timestamped|-ONNX)$/, "");
}

describe("MODELS catalog", () => {
  test("contains the curated D-006 lineup from the onnx-community org", () => {
    const ids = MODELS.map((model) => model.id);
    expect(ids).toEqual([
      "onnx-community/whisper-tiny.en",
      "onnx-community/whisper-tiny",
      "onnx-community/whisper-base.en",
      "onnx-community/whisper-base",
      "onnx-community/whisper-small.en",
      "onnx-community/whisper-small",
      "onnx-community/whisper-medium.en_timestamped",
      "onnx-community/whisper-medium-ONNX",
      "onnx-community/whisper-large-v3-turbo",
    ]);
    expect(new Set(ids).size).toBe(ids.length);
  });

  test(".en variants are English-only; the rest are multilingual", () => {
    for (const model of MODELS) {
      // Hub repo names may carry decorations (e.g. "_timestamped", "-ONNX")
      // after the size suffix — judge English-ness on the undecorated name.
      expect(model.multilingual).toBe(!stripRepoDecorations(model.id).endsWith(".en"));
    }
  });

  test("every model has dtype defaults with notes for both devices", () => {
    for (const model of MODELS) {
      for (const device of ["wasm", "webgpu"] as const) {
        const entry = model.defaults[device];
        expect(entry.dtype).toBeDefined();
        expect(entry.note.length).toBeGreaterThan(0);
        expect(typeof entry.verified).toBe("boolean");
      }
    }
  });

  test("only the M1-tested combination is marked verified", () => {
    const verified = MODELS.flatMap((model) =>
      (["wasm", "webgpu"] as const)
        .filter((device) => model.defaults[device].verified)
        .map((device) => `${model.id}@${device}`),
    );
    expect(verified).toEqual(["onnx-community/whisper-tiny.en@wasm"]);
  });

  test("the default model exists in the catalog", () => {
    expect(getModel(DEFAULT_MODEL_ID)).toBeDefined();
    expect(getModel("onnx-community/whisper-nonexistent")).toBeUndefined();
  });

  test("every model maps to a faster-whisper name for the server engine (D-013)", () => {
    for (const model of MODELS) {
      // The onnx-community ids and faster-whisper's aliases share the same
      // size suffix ("tiny.en", …, "large-v3-turbo") once repo-name
      // decorations ("_timestamped", "-ONNX") are stripped; keeping them
      // aligned is what lets one catalog drive both engines.
      expect(model.fasterWhisperModel).toBe(
        stripRepoDecorations(model.id).replace("onnx-community/whisper-", ""),
      );
    }
    expect(getFasterWhisperModel(DEFAULT_MODEL_ID)).toBe("tiny.en");
    // Unknown ids stay undefined so /api/transcribe can answer 400 instead of
    // forwarding arbitrary strings to the sidecar.
    expect(getFasterWhisperModel("someone/some-model")).toBeUndefined();
  });
});

describe("resolveDefaultDtype", () => {
  test("WASM defaults: q8 for small models, q4 for large ones", () => {
    expect(resolveDefaultDtype("onnx-community/whisper-tiny.en", "wasm")).toBe("q8");
    expect(resolveDefaultDtype("onnx-community/whisper-small", "wasm")).toBe("q8");
    expect(resolveDefaultDtype("onnx-community/whisper-medium-ONNX", "wasm")).toBe("q4");
    expect(resolveDefaultDtype("onnx-community/whisper-large-v3-turbo", "wasm")).toBe("q4");
  });

  test("WebGPU defaults: fp32 for tiny/base (fp16 known-bad), fp16+q4 above", () => {
    expect(resolveDefaultDtype("onnx-community/whisper-tiny", "webgpu")).toEqual({
      encoder_model: "fp32",
      decoder_model_merged: "fp32",
    });
    expect(resolveDefaultDtype("onnx-community/whisper-base.en", "webgpu")).toEqual({
      encoder_model: "fp32",
      decoder_model_merged: "fp32",
    });
    expect(
      resolveDefaultDtype("onnx-community/whisper-large-v3-turbo", "webgpu"),
    ).toEqual({ encoder_model: "fp16", decoder_model_merged: "q4" });
  });

  test("medium never defaults to fp16 on any backend (D-022: broken conversions)", () => {
    for (const id of [
      "onnx-community/whisper-medium.en_timestamped",
      "onnx-community/whisper-medium-ONNX",
    ]) {
      for (const device of ["wasm", "webgpu", "webnn"] as const) {
        expect(resolveDefaultDtype(id, device)).toBe("q4");
      }
    }
  });

  test("unknown model ids fall back to q8", () => {
    expect(resolveDefaultDtype("someone/some-model", "wasm")).toBe("q8");
  });
});

describe("size estimates", () => {
  const tiny = getModel("onnx-community/whisper-tiny.en")!;
  const turbo = getModel("onnx-community/whisper-large-v3-turbo")!;

  test("scale with quantization level", () => {
    const fp32 = estimateDownloadBytes(tiny, "fp32");
    const fp16 = estimateDownloadBytes(tiny, "fp16");
    const q8 = estimateDownloadBytes(tiny, "q8");
    const q4 = estimateDownloadBytes(tiny, "q4");
    expect(fp32).toBeGreaterThan(fp16);
    expect(fp16).toBeGreaterThan(q8);
    expect(q8).toBeGreaterThan(q4);
  });

  test("per-part dtypes are priced per part", () => {
    const mixed = estimateDownloadBytes(turbo, {
      encoder_model: "fp16",
      decoder_model_merged: "q4",
    });
    expect(mixed).toBeLessThan(estimateDownloadBytes(turbo, "fp16"));
    expect(mixed).toBeGreaterThan(estimateDownloadBytes(turbo, "q4"));
  });

  test("default-dtype estimates are in a plausible range", () => {
    // tiny.en on WASM (q8, ~39M params) — tens of MB.
    const tinyWasm = estimateDefaultDownloadBytes(tiny, "wasm");
    expect(tinyWasm).toBeGreaterThan(20e6);
    expect(tinyWasm).toBeLessThan(100e6);
    // large-v3-turbo on WebGPU (fp16 encoder + q4 decoder) — around 1.4 GB.
    const turboWebgpu = estimateDefaultDownloadBytes(turbo, "webgpu");
    expect(turboWebgpu).toBeGreaterThan(1e9);
    expect(turboWebgpu).toBeLessThan(2e9);
  });

  test("formatModelSize renders MB below 1 GB and GB above", () => {
    expect(formatModelSize(42_000_000)).toBe("42 MB");
    expect(formatModelSize(1_400_000_000)).toBe("1.4 GB");
  });
});

describe("estimatePeakLoadBytes (D-027)", () => {
  const medium = getModel("onnx-community/whisper-medium-ONNX")!;
  const turbo = getModel("onnx-community/whisper-large-v3-turbo")!;

  test("is ~2x the larger of the two weight files, not their sum", () => {
    const dtype = resolveDefaultDtype(medium.id, "wasm"); // q4
    const peak = estimatePeakLoadBytes(medium, dtype);
    const download = estimateDownloadBytes(medium, dtype);
    // The combined download sums both files; the peak only doubles the
    // larger one — smaller than 2x the sum, but bigger than the sum alone
    // once the sequential-load patch keeps only one file resident at a time.
    expect(peak).toBeLessThan(2 * download);
    expect(peak).toBeGreaterThan(download);
  });

  test("mixed per-part dtypes (fp16 encoder + q4 decoder) peak on the larger encoder", () => {
    const dtype = resolveDefaultDtype(turbo.id, "webgpu"); // fp16 encoder, q4 decoder
    const peak = estimatePeakLoadBytes(turbo, dtype);
    // large-v3-turbo's encoder (635M) dwarfs its decoder (174M) even before
    // fp16 vs q4 is factored in, so the encoder alone should set the peak.
    const encoderOnlyPeak = 2 * turbo.encoderParamsM * 1e6 * 2; // fp16 = 2 bytes/param
    expect(peak).toBeCloseTo(encoderOnlyPeak, -5);
  });

  test("scales up with a larger model at the same dtype", () => {
    const tiny = getModel("onnx-community/whisper-tiny.en")!;
    expect(estimatePeakLoadBytes(medium, "q4")).toBeGreaterThan(
      estimatePeakLoadBytes(tiny, "q4"),
    );
  });
});

describe("WHISPER_LANGUAGES", () => {
  test("codes are unique, lowercase, and include common languages", () => {
    const codes = WHISPER_LANGUAGES.map((entry) => entry.code);
    expect(new Set(codes).size).toBe(codes.length);
    for (const code of codes) expect(code).toMatch(/^[a-z]{2,3}$/);
    expect(codes).toContain("en");
    expect(codes).toContain("da");
  });

  test("does not contain the auto-detect sentinel (UI adds it)", () => {
    expect(WHISPER_LANGUAGES.some((entry) => entry.code === "auto")).toBe(false);
  });
});

describe("canTranslate (D-WS4a)", () => {
  test("Whisper translate is offered only where OpenAI trained it", () => {
    const translating = MODELS.filter((model) => model.canTranslate).map((model) => model.id);
    // large-v3-turbo was not trained for translation (OpenAI model card);
    // `.en` models are English-only and have no task token.
    expect(translating).toEqual([
      "onnx-community/whisper-tiny",
      "onnx-community/whisper-base",
      "onnx-community/whisper-small",
      "onnx-community/whisper-medium-ONNX",
    ]);
  });

  test("canTranslate implies multilingual", () => {
    for (const model of MODELS) {
      if (model.canTranslate) expect(model.multilingual).toBe(true);
    }
  });
});

describe("taskForModel (D-WS4a)", () => {
  test("a stale translate resets to transcribe on a model that cannot translate", () => {
    expect(taskForModel("onnx-community/whisper-large-v3-turbo", "translate")).toBe("transcribe");
    expect(taskForModel("onnx-community/whisper-small.en", "translate")).toBe("transcribe");
  });

  test("translate survives on a model that can translate", () => {
    expect(taskForModel("onnx-community/whisper-small", "translate")).toBe("translate");
    expect(taskForModel("onnx-community/whisper-small", "transcribe")).toBe("transcribe");
  });

  test("unknown ids never translate", () => {
    expect(taskForModel("someone/unknown", "translate")).toBe("transcribe");
  });
});