import { describe, expect, test } from "bun:test";
import { fetchModelFile, type ModelFileCacheStorage } from "../src/worker/modelFiles";
import type { LoadProgressMessage } from "../src/worker/protocol";

// Cache-API-backed weight download for the Parakeet path (D-WS7a, D-005).

const URL_A = "https://huggingface.co/org/model/resolve/rev/encoder.onnx";

function fakeCaches(initial: Record<string, Uint8Array> = {}) {
  const stored = new Map<string, Uint8Array>(Object.entries(initial));
  const opened: string[] = [];
  const storage: ModelFileCacheStorage = {
    async open(name) {
      opened.push(name);
      return {
        async match(url) {
          const bytes = stored.get(url);
          return bytes ? new Response(bytes as Uint8Array<ArrayBuffer>) : undefined;
        },
        async put(url, response) {
          stored.set(url, new Uint8Array(await response.arrayBuffer()));
        },
      };
    },
  };
  return { storage, stored, opened };
}

function streamingFetch(bytes: Uint8Array, pieces: number, status = 200) {
  const calls: string[] = [];
  const fetchFn = async (url: string) => {
    calls.push(url);
    if (status !== 200) return new Response("nope", { status });
    const size = Math.ceil(bytes.length / pieces);
    const body = new ReadableStream<Uint8Array>({
      start(controller) {
        for (let i = 0; i < bytes.length; i += size) controller.enqueue(bytes.slice(i, i + size));
        controller.close();
      },
    });
    return new Response(body, { headers: { "Content-Length": String(bytes.length) } });
  };
  return { fetchFn, calls };
}

describe("fetchModelFile", () => {
  test("a cached file is returned without touching the network", async () => {
    const cached = new Uint8Array([1, 2, 3]);
    const { storage, opened } = fakeCaches({ [URL_A]: cached });
    const { fetchFn, calls } = streamingFetch(new Uint8Array(), 1);
    const events: LoadProgressMessage[] = [];
    const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
      caches: storage,
      fetch: fetchFn,
      onProgress: (event) => events.push(event),
    });
    expect([...bytes]).toEqual([1, 2, 3]);
    expect(calls).toEqual([]);
    expect(opened).toEqual(["transformers-cache"]);
    expect(events.map((event) => event.status)).toEqual(["initiate", "done"]);
  });

  test("a miss downloads with byte progress and stores the complete file", async () => {
    const payload = new Uint8Array(1000).map((_, i) => i % 251);
    const { storage, stored } = fakeCaches();
    const { fetchFn, calls } = streamingFetch(payload, 4);
    const events: LoadProgressMessage[] = [];
    const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
      caches: storage,
      fetch: fetchFn,
      onProgress: (event) => events.push(event),
    });
    expect(calls).toEqual([URL_A]);
    expect([...bytes]).toEqual([...payload]);
    expect([...stored.get(URL_A)!]).toEqual([...payload]);
    const progress = events.filter((event) => event.status === "progress");
    expect(progress.length).toBe(4);
    const last = progress[progress.length - 1]!;
    expect(last).toMatchObject({ file: "encoder.onnx", loaded: 1000, total: 1000, progress: 100 });
    expect(events[0]!.status).toBe("initiate");
    expect(events[events.length - 1]!.status).toBe("done");
  });

  test("an HTTP error throws a plain message and stores nothing", async () => {
    const { storage, stored } = fakeCaches();
    const { fetchFn } = streamingFetch(new Uint8Array(), 1, 404);
    await expect(
      fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }),
    ).rejects.toThrow("Could not download encoder.onnx (HTTP 404)");
    expect(stored.size).toBe(0);
  });

  test("an interrupted download stores nothing (no false 'available offline')", async () => {
    const { storage, stored } = fakeCaches();
    const fetchFn = async () =>
      new Response(
        new ReadableStream<Uint8Array>({
          start(controller) {
            controller.enqueue(new Uint8Array(10));
            controller.error(new Error("network down"));
          },
        }),
        { headers: { "Content-Length": "100" } },
      );
    await expect(
      fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }),
    ).rejects.toThrow();
    expect(stored.size).toBe(0);
  });

  test("without Cache Storage it still downloads", async () => {
    const { fetchFn, calls } = streamingFetch(new Uint8Array([9]), 1);
    const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
      caches: undefined,
      fetch: fetchFn,
      onProgress: () => {},
    });
    expect([...bytes]).toEqual([9]);
    expect(calls).toEqual([URL_A]);
  });
});