import sharp from "sharp";
import { describe, expect, it } from "vitest";
import type { ComfyUiProviderConfig } from "./comfyui-provider.js";
import {
  createComfyUiProvider,
  describeExecutionError,
  parseGpuStatus,
} from "./comfyui-provider.js";
import { ComfyUiEmulator } from "./testing/comfyui-emulator.js";

const BASE_URL = "http://127.0.0.1:8188";
const IMAGE = { prompt: "a lighthouse", width: 512, height: 512 };

function setup(config: Partial<ComfyUiProviderConfig> = {}, emulator = new ComfyUiEmulator()) {
  const provider = createComfyUiProvider({
    baseUrl: BASE_URL,
    fetch: emulator.fetch,
    webSocket: emulator.openWebSocket,
    pollIntervalMs: 1,
    ...config,
  });
  return { provider, emulator };
}

/** Progress rounded so float sums compare exactly. */
function rounded(progress: number[]): number[] {
  return progress.map((fraction) => Number(fraction.toFixed(4)));
}

function submitted(emulator: ComfyUiEmulator, index = 0) {
  const graph = emulator.prompts[index]?.graph ?? {};
  const byClass = (classType: string) =>
    Object.values(graph).find((node) => node.class_type === classType)?.inputs ?? {};
  return { graph, byClass };
}

describe("createComfyUiProvider", () => {
  it("refuses metadata hosts, public hosts and keys over public http", () => {
    expect(() => createComfyUiProvider({ baseUrl: "http://169.254.169.254" })).toThrow(
      "blocked metadata or link-local host",
    );
    expect(() => createComfyUiProvider({ baseUrl: "http://comfy.example.com" })).toThrow(
      "Public model endpoints are blocked",
    );
    expect(() => createComfyUiProvider({ baseUrl: "ftp://127.0.0.1" })).toThrow("http or https");
  });

  it("validates custom workflows up front", () => {
    expect(() => createComfyUiProvider({ baseUrl: BASE_URL, imageWorkflow: "{" })).toThrow(
      "Workflow is not valid JSON",
    );
    expect(() =>
      createComfyUiProvider({ baseUrl: BASE_URL, videoWorkflow: '{"nodes": [], "links": []}' }),
    ).toThrow("API format");
  });

  it("keeps a reverse-proxy path prefix and sends the API key", async () => {
    const emulator = new ComfyUiEmulator({ apiKey: "comfy-key", basePath: "/comfy" });
    const { provider } = setup({ baseUrl: `${BASE_URL}/comfy/`, apiKey: "comfy-key" }, emulator);
    await provider.generateImage(IMAGE);
    expect(emulator.requests.map((request) => request.path)).toContain("prompt");
    expect(emulator.requests.every((request) => request.authorization === "Bearer comfy-key")).toBe(
      true,
    );
    const { provider: keyless } = setup({ baseUrl: `${BASE_URL}/comfy` }, emulator);
    await expect(keyless.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 401");
  });

  it("fills the default text-to-image workflow", async () => {
    const { provider, emulator } = setup();
    await provider.generateImage({ ...IMAGE, width: 1001, negativePrompt: "fog", seed: 42 });
    const { byClass } = submitted(emulator);
    expect(byClass("CheckpointLoaderSimple").ckpt_name).toBe("sd_xl_base_1.0.safetensors");
    expect(byClass("EmptyLatentImage")).toMatchObject({ width: 1000, height: 512 });
    expect(byClass("KSampler")).toMatchObject({ seed: 42, denoise: 1 });
    const texts = emulator.prompts[0]!.graph;
    expect(Object.values(texts).map((node) => node.inputs.text)).toEqual(
      expect.arrayContaining(["a lighthouse", "fog"]),
    );
  });

  it("uses the configured model, else the first checkpoint the video workflows do not use", async () => {
    const emulator = new ComfyUiEmulator();
    emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors", "v1-5-pruned.safetensors"];
    emulator.checkpointListFormat = "legacy";
    await setup({}, emulator).provider.generateImage(IMAGE);
    expect(submitted(emulator, 0).byClass("CheckpointLoaderSimple").ckpt_name).toBe(
      "v1-5-pruned.safetensors",
    );
    emulator.checkpoints.push("dreamshaper.safetensors");
    await setup({ imageModel: "dreamshaper.safetensors" }, emulator).provider.generateImage(IMAGE);
    expect(submitted(emulator, 1).byClass("CheckpointLoaderSimple").ckpt_name).toBe(
      "dreamshaper.safetensors",
    );
    emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors"];
    await expect(setup({}, emulator).provider.generateImage(IMAGE)).rejects.toThrow(
      "ComfyUI has no image checkpoint installed",
    );
  });

  it("uploads the source image for image-to-image", async () => {
    const { provider, emulator } = setup();
    const bytes = new Uint8Array(Buffer.from((await sourcePng()).bytes));
    await provider.generateImage({
      ...IMAGE,
      source: { bytes, mimeType: "image/png" },
      strength: 0.3,
    });
    const { byClass } = submitted(emulator);
    const uploaded = String(byClass("LoadImage").image);
    expect(uploaded).toMatch(/^evoo-.+\.png$/);
    expect(Buffer.from(emulator.uploads.get(uploaded)!).equals(Buffer.from(bytes))).toBe(true);
    expect(byClass("KSampler").denoise).toBe(0.3);
  });

  it("lets a custom workflow replace only the default it matches", async () => {
    const custom = {
      "1": {
        class_type: "CustomTextToImage",
        inputs: { text: "style: {{prompt}}", seed: "{{seed}}" },
      },
      "2": { class_type: "SaveImage", inputs: { images: ["1", 0], filename_prefix: "x" } },
    };
    const { provider, emulator } = setup({ imageWorkflow: JSON.stringify(custom) });
    await provider.generateImage(IMAGE);
    expect(submitted(emulator, 0).byClass("CustomTextToImage").text).toBe("style: a lighthouse");
    await provider.generateImage({ ...IMAGE, source: await sourcePng() });
    expect(submitted(emulator, 1).byClass("LoadImage").image).toBeDefined();
  });

  it("fills the LTX-Video workflow with an 8n+1 frame count and reports the clip", async () => {
    const { provider, emulator } = setup();
    const [clip] = await provider.generateVideo({
      prompt: "a paper boat",
      width: 470,
      height: 836,
      durationSeconds: 5,
    });
    const { byClass } = submitted(emulator);
    expect(byClass("EmptyLTXVLatentVideo")).toMatchObject({ width: 480, height: 832, length: 121 });
    expect(byClass("CreateVideo").fps).toBe(24);
    expect(clip).toMatchObject({
      mimeType: "video/mp4",
      width: 480,
      height: 832,
      durationMs: 5042,
    });
  });

  it("reads video from a custom workflow's gifs output and skips frame previews", async () => {
    const custom = {
      "1": { class_type: "TextToFrames", inputs: { text: "{{prompt}}" } },
      "2": { class_type: "PreviewImage", inputs: { images: ["1", 0] } },
      "3": { class_type: "VHS_VideoCombine", inputs: { images: ["1", 0], frame_rate: 16 } },
    };
    const { provider } = setup({ videoWorkflow: JSON.stringify(custom) });
    const clips = await provider.generateVideo({ ...IMAGE, durationSeconds: 2 });
    expect(clips.map((clip) => [clip.mimeType, clip.durationMs])).toEqual([
      ["video/mp4", undefined],
    ]);
  });

  it("explains rejected workflows and failed runs", async () => {
    const emulator = new ComfyUiEmulator();
    const { provider } = setup({ imageModel: "missing.safetensors" }, emulator);
    await expect(provider.generateImage(IMAGE)).rejects.toThrow(
      "ComfyUI rejected the workflow: CheckpointLoaderSimple: Value not in list: ckpt_name: 'missing.safetensors' not in list",
    );
    const working = setup({}, emulator).provider;
    emulator.failNextExecution = "CUDA out of memory. Tried to allocate 2.00 GiB\nTraceback...";
    await expect(working.generateImage(IMAGE)).rejects.toThrow(
      /^ComfyUI failed in KSampler: CUDA out of memory\. Tried to allocate 2\.00 GiB$/,
    );
    emulator.failNextStatus = 500;
    await expect(working.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 500");
  });

  it("reports an unreachable server", async () => {
    const provider = createComfyUiProvider({
      baseUrl: BASE_URL,
      fetch: async () => {
        throw new TypeError("fetch failed");
      },
    });
    await expect(provider.generateImage(IMAGE)).rejects.toThrow("ComfyUI is unreachable");
    await expect(provider.gpuStatus()).rejects.toThrow("ComfyUI is unreachable");
  });

  it("interrupts its own running prompt when the caller cancels", async () => {
    const emulator = new ComfyUiEmulator();
    emulator.pendingPolls = 1_000_000;
    const { provider } = setup({ pollIntervalMs: 5 }, emulator);
    const controller = new AbortController();
    const run = provider.generateImage({ ...IMAGE, signal: controller.signal });
    while (!emulator.requests.some((request) => request.path.startsWith("history/"))) {
      await new Promise((resolve) => setTimeout(resolve, 1));
    }
    controller.abort(new Error("cancelled by user"));
    await expect(run).rejects.toThrow("cancelled by user");
    expect(emulator.deleted).toEqual(["prompt-1"]);
    expect(emulator.interrupted).toEqual(["prompt-1"]);
  });

  it("reports live step progress in order, short of done until the outputs are in", async () => {
    const { provider, emulator } = setup();
    emulator.pendingPolls = 4;
    const progress: number[] = [];
    await provider.generateImage({ ...IMAGE, onProgress: (fraction) => progress.push(fraction) });
    expect(progress).toEqual([0.25, 0.5, 0.75, 0.99, 1]);
    const [socket] = emulator.sockets;
    expect(socket?.url.href).toBe(
      `ws://127.0.0.1:8188/ws?clientId=${emulator.prompts[0]?.clientId}`,
    );
    expect(socket?.closedByClient).toBe(true);
  });

  it("maps progress across the prompts of a job", async () => {
    const { provider, emulator } = setup();
    emulator.pendingPolls = 2;
    const progress: number[] = [];
    await provider.generateImage({
      ...IMAGE,
      count: 2,
      onProgress: (fraction) => progress.push(fraction),
    });
    expect(rounded(progress)).toEqual([0.25, 0.495, 0.5, 0.75, 0.995, 1]);
    expect(emulator.sockets).toHaveLength(1);
  });

  it("sends the key and keeps the path prefix on the websocket", async () => {
    const emulator = new ComfyUiEmulator({ apiKey: "comfy-key", basePath: "/comfy" });
    const { provider } = setup(
      { baseUrl: "https://127.0.0.1:8188/comfy", apiKey: "comfy-key" },
      emulator,
    );
    const progress: number[] = [];
    await provider.generateImage({ ...IMAGE, onProgress: (fraction) => progress.push(fraction) });
    const socket = emulator.sockets[0]!;
    expect(`${socket.url.origin}${socket.url.pathname}`).toBe("wss://127.0.0.1:8188/comfy/ws");
    expect(socket.authorization).toBe("Bearer comfy-key");
    expect(progress).toEqual([0.99, 1]);
  });

  it("generates as before when the websocket is unavailable", async () => {
    const emulator = new ComfyUiEmulator();
    emulator.webSocketAvailable = false;
    const progress: number[] = [];
    const onProgress = (fraction: number) => progress.push(fraction);
    await setup({}, emulator).provider.generateImage({ ...IMAGE, onProgress });
    const throwing = setup({
      webSocket: () => {
        throw new Error("no websocket");
      },
    });
    await throwing.provider.generateImage({ ...IMAGE, onProgress });
    // A substituted fetch without a matching socket opens none, rather than a real connection.
    const fetchOnly = new ComfyUiEmulator();
    await setup({ webSocket: undefined }, fetchOnly).provider.generateImage({
      ...IMAGE,
      onProgress,
    });
    expect(progress).toEqual([1, 1, 1]);
    expect(emulator.sockets.map((socket) => socket.readyState)).toEqual([3]);
    expect(fetchOnly.sockets).toEqual([]);
  });

  it("ignores stray messages, a socket that drops mid-run and a failing listener", async () => {
    const { provider, emulator } = setup();
    emulator.pendingPolls = 3;
    const progress: number[] = [];
    await provider.generateImage({
      ...IMAGE,
      onProgress: (fraction) => {
        progress.push(fraction);
        const socket = emulator.sockets[0]!;
        socket.push("not json");
        socket.push({ type: "progress", data: { value: "2", max: 3, prompt_id: "prompt-1" } });
        socket.push({ type: "progress", data: { value: 2, max: 0, prompt_id: "prompt-1" } });
        socket.push({ type: "progress", data: { value: 3, max: 3, prompt_id: "prompt-9" } });
        socket.drop();
      },
    });
    expect(rounded(progress)).toEqual([0.3333, 1]);
    const failing = setup().provider;
    await expect(
      failing.generateImage({
        ...IMAGE,
        onProgress: () => {
          throw new Error("listener failed");
        },
      }),
    ).resolves.toHaveLength(1);
  });

  it("closes the websocket when the caller cancels or the run fails", async () => {
    const emulator = new ComfyUiEmulator();
    emulator.pendingPolls = 1_000_000;
    const { provider } = setup({ pollIntervalMs: 5 }, emulator);
    const controller = new AbortController();
    const onProgress = () => undefined;
    const run = provider.generateImage({ ...IMAGE, signal: controller.signal, onProgress });
    while (!emulator.requests.some((request) => request.path.startsWith("history/"))) {
      await new Promise((resolve) => setTimeout(resolve, 1));
    }
    expect(emulator.sockets[0]?.readyState).toBe(1);
    controller.abort(new Error("cancelled by user"));
    expect(emulator.sockets[0]?.closedByClient).toBe(true);
    await expect(run).rejects.toThrow("cancelled by user");
    emulator.pendingPolls = 1;
    emulator.failNextExecution = "CUDA out of memory";
    await expect(provider.generateImage({ ...IMAGE, onProgress })).rejects.toThrow(
      "CUDA out of memory",
    );
    expect(emulator.sockets.map((socket) => socket.readyState)).toEqual([3, 3]);
  });

  it("reports GPU memory other processes leave to ComfyUI", async () => {
    const emulator = new ComfyUiEmulator();
    emulator.gpu = { totalMb: 12_288, freeMb: 3_000, torchReservedMb: 7_000, torchFreeMb: 1_000 };
    await expect(setup({}, emulator).provider.gpuStatus()).resolves.toEqual({
      vramFreeMb: 9_000,
      vramTotalMb: 12_288,
    });
    emulator.gpu = null;
    await expect(setup({}, emulator).provider.gpuStatus()).resolves.toBeUndefined();
    expect(parseGpuStatus({ devices: [{ type: "cuda", vram_total: 0, vram_free: 0 }] })).toBe(
      undefined,
    );
  });

  it("names interrupted runs", () => {
    expect(describeExecutionError([["execution_interrupted", {}]])).toBe(
      "ComfyUI interrupted the run",
    );
    expect(describeExecutionError(undefined)).toBe("ComfyUI failed");
  });
});

async function sourcePng() {
  const png = await sharp({
    create: { width: 32, height: 32, channels: 3, background: { r: 1, g: 2, b: 3 } },
  })
    .png()
    .toBuffer();
  return { bytes: new Uint8Array(png), mimeType: "image/png" };
}