import { describe, expect, test } from "bun:test";
import { buildSdxlWorkflow, selectGeneratedOutputImage } from "./comfy";

describe("Comfy adapter", () => {
  test("selects SaveImage output instead of uploaded input or mask images", () => {
    const image = selectGeneratedOutputImage({
      outputs: {
        "4": { images: [{ filename: "image-studio-input.png", type: "input" }] },
        "9": { images: [{ filename: "image-studio-mask.png", type: "input" }] },
        "8": { images: [{ filename: "image-studio-inpaint_00001_.png", subfolder: "", type: "output" }] },
      },
    });

    expect(image).toEqual({ filename: "image-studio-inpaint_00001_.png", subfolder: "", type: "output" });
  });

  test("falls back to generated filename prefixes when node ids differ", () => {
    const image = selectGeneratedOutputImage({
      outputs: {
        "12": { images: [{ filename: "image-studio-inpaint_00002_.png", type: "output" }] },
        "4": { images: [{ filename: "image-studio-input.png", type: "input" }] },
      },
    });

    expect(image?.filename).toBe("image-studio-inpaint_00002_.png");
  });

  test("builds neutral inpaint with VAEEncodeForInpaint", () => {
    const workflow = buildSdxlWorkflow(inpaintRequest({ maskedContent: "neutral" }));

    expect(workflow["5"]?.class_type).toBe("VAEEncodeForInpaint");
    expect(workflow["5"]?.inputs).toMatchObject({ grow_mask_by: 6, mask: ["11", 0] });
    expect(workflow["6"]?.inputs.latent_image).toEqual(["5", 0]);
  });

  test("builds original-content inpaint with a latent noise mask", () => {
    const workflow = buildSdxlWorkflow(inpaintRequest({ maskedContent: "original", growMaskBy: 12 }));

    expect(workflow["5"]?.class_type).toBe("VAEEncode");
    expect(workflow["12"]?.class_type).toBe("GrowMask");
    expect(workflow["12"]?.inputs).toMatchObject({ mask: ["11", 0], expand: 12 });
    expect(workflow["13"]?.class_type).toBe("SetLatentNoiseMask");
    expect(workflow["13"]?.inputs).toMatchObject({ samples: ["5", 0], mask: ["12", 0] });
    expect(workflow["6"]?.inputs.latent_image).toEqual(["13", 0]);
  });
});

function inpaintRequest(inpaint: { maskedContent: "neutral" | "original"; growMaskBy?: number }) {
  return {
    mode: "inpaint" as const,
    model: "model.safetensors",
    prompt: "replace garment",
    width: 128,
    height: 128,
    inputImage: "input.png",
    maskImage: "mask.png",
    inpaint,
  };
}