type GenerateArchitecture = "sdxl" | "z-image" | "z-image-turbo" | "anima";
type GenerateMode = "text-to-image" | "image-to-image" | "inpaint" | "outpaint";
type Workflow = Record<string, { class_type: string; inputs: Record<string, unknown> }>;

export type ComfyGenerateRequest = {
  architecture?: GenerateArchitecture;
  mode: GenerateMode;
  model: string;
  textEncoder?: string;
  vae?: string;
  prompt: string;
  negativePrompt?: string;
  strength?: number;
  steps?: number;
  cfg?: number;
  seed?: number;
  sampler?: string;
  scheduler?: string;
  width?: number;
  height?: number;
  batchSize?: number;
  refinePass?: boolean;
  refineStrength?: number;
  outpaint?: {
    left?: number;
    top?: number;
    right?: number;
    bottom?: number;
    feathering?: number;
  };
  inpaint?: {
    growMaskBy?: number;
    maskBlur?: number;
    maskFeather?: number;
    maskExpand?: number;
    cropPadding?: number;
    maskPolarity?: "hidden" | "revealed";
    maskedContent?: "neutral" | "original" | "originalColor" | "edges";
    crop?: unknown;
    placement?: unknown;
    structureControl?: "none" | "canny" | "depth" | "pose";
    controlStrength?: number;
    controlModel?: string;
  };
  inputImage?: string;
  maskImage?: string;
};

export type ComfyProgress = { progress: number; detail: string };
export type ComfySegmentRequest = { inputImage: string; x: number; y: number; model?: string };

type ComfyObjectInfo = {
  CheckpointLoaderSimple?: { input?: { required?: { ckpt_name?: [string[]] } } };
  KSampler?: { input?: { required?: { sampler_name?: [string[]]; scheduler?: [string[]] } } };
  UNETLoader?: { input?: { required?: { unet_name?: [string[]] } } };
  CLIPLoader?: { input?: { required?: { clip_name?: [string[]] } } };
  VAELoader?: { input?: { required?: { vae_name?: [string[]] } } };
  ControlNetLoader?: { input?: { required?: { control_net_name?: [string[]] } } };
  Canny?: unknown;
  "MiDaS-DepthMapPreprocessor"?: unknown;
  OpenposePreprocessor?: unknown;
  ControlNetApplyAdvanced?: unknown;
  SAM3_Detect?: unknown;
  MaskToImage?: unknown;
};

const comfyBaseUrl = process.env.COMFYUI_URL ?? "http://127.0.0.1:8188";
const comfyHistoryTimeoutMs = parsePositiveInteger(process.env.COMFYUI_HISTORY_TIMEOUT_MS, 20 * 60 * 1000);
const comfyHistoryPollIntervalMs = parsePositiveInteger(process.env.COMFYUI_HISTORY_POLL_INTERVAL_MS, 1000);
const defaultModels: Record<GenerateArchitecture, string> = {
  sdxl: "sd_xl_base_1.0.safetensors",
  "z-image": "z_image_bf16.safetensors",
  "z-image-turbo": "z_image_turbo_bf16.safetensors",
  anima: "anima-base-v1.0.safetensors",
};

export async function listGenerationOptions() {
  const info = await fetchObjectInfo();
  const checkpointModels = info.CheckpointLoaderSimple?.input?.required?.ckpt_name?.[0] ?? [];
  const diffusionModels = info.UNETLoader?.input?.required?.unet_name?.[0] ?? [];
  const textEncoders = info.CLIPLoader?.input?.required?.clip_name?.[0] ?? [];
  const vaes = info.VAELoader?.input?.required?.vae_name?.[0] ?? [];

  return {
    models: checkpointModels,
    inpaintModels: checkpointModels.filter((model) => /inpaint|fill/i.test(model)),
    samplers: info.KSampler?.input?.required?.sampler_name?.[0] ?? [],
    schedulers: info.KSampler?.input?.required?.scheduler?.[0] ?? [],
    diffusionModels,
    textEncoders,
    vaes,
    controlModels: info.ControlNetLoader?.input?.required?.control_net_name?.[0] ?? [],
    structureControls: [
      ...(info.Canny && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["canny"] : []),
      ...(info["MiDaS-DepthMapPreprocessor"] && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["depth"] : []),
      ...(info.OpenposePreprocessor && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["pose"] : []),
    ],
    semanticSelection: Boolean(info.SAM3_Detect && info.MaskToImage && checkpointModels.some((model) => /sam.?3/i.test(model))),
    sam3Models: checkpointModels.filter((model) => /sam.?3/i.test(model)),
    architectures: [
      {
        value: "sdxl",
        label: "SDXL",
        defaultModel: checkpointModels[0] ?? defaultModels.sdxl,
        models: checkpointModels,
        supportedModes: ["text-to-image", "image-to-image", "inpaint", "outpaint"],
      },
      {
        value: "z-image",
        label: "Z-Image",
        defaultModel: defaultModels["z-image"],
        models: modelsForArchitecture(diffusionModels, defaultModels["z-image"], [/z[_-]?image(?!.*turbo)/i]),
        supportedModes: ["text-to-image"],
      },
      {
        value: "z-image-turbo",
        label: "Z-Image Turbo",
        defaultModel: defaultModels["z-image-turbo"],
        models: modelsForArchitecture(diffusionModels, defaultModels["z-image-turbo"], [/z[_-]?image.*turbo/i]),
        supportedModes: ["text-to-image"],
      },
      {
        value: "anima",
        label: "Anima",
        defaultModel: defaultModels.anima,
        models: modelsForArchitecture(diffusionModels, defaultModels.anima, [], { exclude: [/z[_-]?image/i] }),
        supportedModes: ["text-to-image"],
      },
    ],
  };
}

export async function segment(request: ComfySegmentRequest, signal?: AbortSignal) {
  if (!request.inputImage) throw new Error("Semantic selection requires an image");
  if (!Number.isFinite(request.x) || !Number.isFinite(request.y)) throw new Error("Semantic selection requires a valid point");
  const info = await fetchObjectInfo();
  const models = info.CheckpointLoaderSimple?.input?.required?.ckpt_name?.[0] ?? [];
  const model = request.model && request.model !== "auto" ? request.model : models.find((candidate) => /sam.?3/i.test(candidate));
  if (!info.SAM3_Detect || !info.MaskToImage || !model || !models.includes(model)) throw new Error("SAM3 semantic selection is not installed in ComfyUI. Install a SAM3 model and enable the native SAM3 nodes.");
  const uploaded = await uploadDataUrl(request.inputImage, `image-studio-segment-${crypto.randomUUID()}.png`, signal);
  const prompt = buildSemanticSelectionWorkflow({ ...request, model, inputImage: uploaded });
  const queued = await fetch(`${comfyBaseUrl}/prompt`, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ client_id: crypto.randomUUID(), prompt }), signal });
  if (!queued.ok) throw new Error(`ComfyUI semantic selection failed: ${queued.status} ${await queued.text()}`);
  const queuedBody = await queued.json() as { prompt_id?: string; node_errors?: unknown };
  const nodeError = nodeErrorsMessage(queuedBody.node_errors);
  if (nodeError) throw new Error(`ComfyUI rejected semantic selection: ${nodeError}`);
  if (!queuedBody.prompt_id) throw new Error("ComfyUI did not return a semantic selection prompt id");
  const history = await waitForHistory(queuedBody.prompt_id, signal);
  const historyError = historyErrorMessage(history);
  if (historyError) throw new Error(`ComfyUI semantic selection failed: ${historyError}`);
  const image = selectGeneratedOutputImage(history);
  if (!image) throw new Error("ComfyUI did not return a semantic selection mask");
  const response = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`, { signal });
  if (!response.ok) throw new Error(`ComfyUI mask fetch failed: ${response.status}`);
  const bytes = Buffer.from(await response.arrayBuffer());
  return { source: `data:image/png;base64,${bytes.toString("base64")}`, mimeType: "image/png" };
}

export function buildSemanticSelectionWorkflow(request: ComfySegmentRequest & { model: string }): Workflow {
  return {
    "1": { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: request.model } },
    "2": { class_type: "LoadImage", inputs: { image: request.inputImage } },
    "3": { class_type: "SAM3_Detect", inputs: { model: ["1", 0], image: ["2", 0], positive_coords: JSON.stringify([{ x: Math.round(request.x), y: Math.round(request.y) }]), threshold: 0.5, refine_iterations: 2, individual_masks: false } },
    "4": { class_type: "MaskToImage", inputs: { mask: ["3", 0] } },
    "8": { class_type: "SaveImage", inputs: { filename_prefix: "image-studio-segment", images: ["4", 0] } },
  };
}

async function listCheckpointModels() {
  return (await listGenerationOptions()).models;
}

export async function generate(request: ComfyGenerateRequest, signal?: AbortSignal, onProgress?: (event: ComfyProgress) => void) {
  if (!request.prompt?.trim()) throw new Error("Prompt is required");
  onProgress?.({ progress: 0.03, detail: "Preparing workflow" });
  const architecture = normalizeArchitecture(request.architecture);
  if (request.mode !== "text-to-image" && architecture !== "sdxl") throw new Error(`${architectureLabel(architecture)} currently supports text-to-image only`);
  if (!request.model || request.model === "auto") request.model = await defaultModelForArchitecture(architecture, request.mode);

  const clientId = crypto.randomUUID();
  const uploaded = request.inputImage ? await uploadDataUrl(request.inputImage, `image-studio-${crypto.randomUUID()}.png`, signal) : undefined;
  const mask = request.maskImage ? await uploadDataUrl(request.maskImage, `image-studio-mask-${crypto.randomUUID()}.png`, signal) : undefined;
  onProgress?.({ progress: 0.16, detail: "Inputs uploaded" });
  if (request.mode === "inpaint" && (!uploaded || !mask)) throw new Error("Inpaint requires normalized input and mask images");
  const resolvedRequest = await resolveStructureControl({ ...request, architecture, inputImage: uploaded, maskImage: mask });
  const prompt = buildComfyWorkflow(resolvedRequest);

  const queued = await fetch(`${comfyBaseUrl}/prompt`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ client_id: clientId, prompt }),
    signal,
  });
  if (!queued.ok) throw new Error(`ComfyUI prompt failed: ${queued.status} ${await queued.text()}`);
  const queuedBody = await queued.json() as { prompt_id?: string; node_errors?: unknown };
  const nodeError = nodeErrorsMessage(queuedBody.node_errors);
  if (nodeError) throw new Error(`ComfyUI rejected the workflow: ${nodeError}`);
  if (!queuedBody.prompt_id) throw new Error("ComfyUI did not return a prompt id");
  onProgress?.({ progress: 0.24, detail: "Queued in ComfyUI" });
  const prompt_id = queuedBody.prompt_id;
  let history: unknown;
  try {
    history = await waitForHistory(prompt_id, signal, onProgress);
  } catch (error) {
    if (signal?.aborted) await cancelComfyPrompt(prompt_id);
    throw error;
  }
  const historyError = historyErrorMessage(history);
  if (historyError) throw new Error(`ComfyUI generation failed: ${historyError}`);
  const images = selectGeneratedOutputImages(history);
  if (images.length === 0) throw new Error("ComfyUI did not return an image");
  onProgress?.({ progress: 0.9, detail: "Downloading results" });
  const results = await Promise.all(images.map(async (image, index) => {
    const imageResponse = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`, { signal });
    if (!imageResponse.ok) throw new Error(`ComfyUI image fetch failed: ${imageResponse.status}`);
    const bytes = Buffer.from(await imageResponse.arrayBuffer());
    return { source: `data:image/png;base64,${bytes.toString("base64")}`, mimeType: "image/png", seed: Math.round((resolvedRequest.seed ?? 0) + index) };
  }));
  onProgress?.({ progress: 1, detail: "Results ready" });
  return { results };
}

async function uploadDataUrl(dataUrl: string, filename: string, signal?: AbortSignal) {
  const match = /^data:([^;]+);base64,(.+)$/.exec(dataUrl);
  if (!match) throw new Error("Expected a base64 data URL image");
  const mimeType = match[1] ?? "image/png";
  const base64 = match[2] ?? "";
  const form = new FormData();
  form.append("image", new File([new Uint8Array(Buffer.from(base64, "base64"))], filename, { type: mimeType }));
  form.append("overwrite", "true");
  const response = await fetch(`${comfyBaseUrl}/upload/image`, { method: "POST", body: form, signal });
  if (!response.ok) throw new Error(`ComfyUI upload failed: ${response.status}`);
  const uploaded = await response.json() as { name: string };
  return uploaded.name;
}

async function waitForHistory(promptId: string, signal?: AbortSignal, onProgress?: (event: ComfyProgress) => void) {
  const startedAt = Date.now();
  let attempts = 0;
  while (Date.now() - startedAt < comfyHistoryTimeoutMs) {
    if (signal?.aborted) {
      await cancelComfyPrompt(promptId);
      throw signal.reason ?? new DOMException("Generation cancelled", "AbortError");
    }
    const response = await fetch(`${comfyBaseUrl}/history/${promptId}`, { signal });
    attempts += 1;
    const elapsedRatio = Math.min(1, (Date.now() - startedAt) / comfyHistoryTimeoutMs);
    onProgress?.({ progress: 0.25 + elapsedRatio * 0.6, detail: attempts <= 1 ? "Generating" : `Generating · check ${attempts}` });
    if (response.ok) {
      const history = await response.json() as Record<string, unknown>;
      if (history[promptId]) return history[promptId];
    }
    await abortableSleep(comfyHistoryPollIntervalMs, signal);
  }
  throw new Error(`Timed out waiting for ComfyUI prompt ${promptId} after ${Math.round(comfyHistoryTimeoutMs / 1000)} seconds and ${attempts} checks`);
}

async function cancelComfyPrompt(promptId: string) {
  const queueResponse = await fetch(`${comfyBaseUrl}/queue`).catch(() => undefined);
  const queue = queueResponse?.ok ? await queueResponse.json() as { queue_running?: unknown[][]; queue_pending?: unknown[][] } : undefined;
  const isRunning = queue?.queue_running?.some((entry) => entry.includes(promptId)) ?? false;
  const isPending = queue?.queue_pending?.some((entry) => entry.includes(promptId)) ?? true;
  const requests: Promise<unknown>[] = [];
  if (isPending) requests.push(fetch(`${comfyBaseUrl}/queue`, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ delete: [promptId] }) }));
  if (isRunning) requests.push(fetch(`${comfyBaseUrl}/interrupt`, { method: "POST" }));
  await Promise.allSettled(requests);
}

function abortableSleep(ms: number, signal?: AbortSignal): Promise<void> {
  if (!signal) return Bun.sleep(ms);
  return new Promise((resolve, reject) => {
    const timeout = setTimeout(resolve, ms);
    signal.addEventListener("abort", () => {
      clearTimeout(timeout);
      reject(signal.reason ?? new DOMException("Generation cancelled", "AbortError"));
    }, { once: true });
  });
}

export function selectGeneratedOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined {
  return selectGeneratedOutputImages(history)[0];
}

export function selectGeneratedOutputImages(history: unknown): Array<{ filename: string; subfolder?: string; type?: string }> {
  const outputs = (history as { outputs?: Record<string, { images?: { filename: string; subfolder?: string; type?: string }[] }> }).outputs ?? {};
  const saveImageOutput = outputs["8"]?.images?.filter(isGeneratedImage) ?? [];
  if (saveImageOutput.length > 0) return saveImageOutput;

  const prefixedOutput = Object.values(outputs).flatMap((output) => output.images ?? []).filter((image) => image.filename.startsWith("image-studio-") && image.type !== "input");
  if (prefixedOutput.length > 0) return prefixedOutput;

  return Object.values(outputs).flatMap((output) => output.images ?? []).filter(isGeneratedImage);
}

function isGeneratedImage(image: { filename: string; subfolder?: string; type?: string }) {
  return image.type === undefined || image.type === "output";
}

export function buildComfyWorkflow(request: ComfyGenerateRequest): Workflow {
  switch (normalizeArchitecture(request.architecture)) {
    case "z-image":
      return buildZImageWorkflow(request);
    case "z-image-turbo":
      return buildZImageTurboWorkflow(request);
    case "anima":
      return buildAnimaWorkflow(request);
    case "sdxl":
    default:
      return buildSdxlWorkflow(request);
  }
}

export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow {
  if (!request.width || !request.height) throw new Error("Generation width and height are required");
  const width = Math.max(64, Math.round(request.width));
  const height = Math.max(64, Math.round(request.height));
  const denoise = Math.max(0, Math.min(1, (request.strength ?? 75) / 100));
  const seed = request.seed === undefined || request.seed < 0 ? Math.floor(Math.random() * 2 ** 32) : Math.round(request.seed);
  const steps = Math.max(1, Math.round(request.steps ?? 30));
  const cfg = Math.max(0, request.cfg ?? 7);
  const sampler = request.sampler ?? "euler";
  const scheduler = request.scheduler ?? "normal";
  const positive = request.prompt;
  const negative = request.negativePrompt ?? "";
  const samplerInputs: Record<string, unknown> = { seed, steps, cfg, sampler_name: sampler, scheduler, denoise, model: ["1", 0], positive: ["2", 0], negative: ["3", 0], latent_image: ["5", 0] };
  const workflow: Workflow = {
    "1": { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: request.model } },
    "2": { class_type: "CLIPTextEncode", inputs: { text: positive, clip: ["1", 1] } },
    "3": { class_type: "CLIPTextEncode", inputs: { text: negative, clip: ["1", 1] } },
    "6": { class_type: "KSampler", inputs: samplerInputs },
    "7": { class_type: "VAEDecode", inputs: { samples: ["6", 0], vae: ["1", 2] } },
    "8": { class_type: "SaveImage", inputs: { filename_prefix: `image-studio-${request.mode}`, images: ["7", 0] } },
  };

  if (request.mode === "text-to-image" || !request.inputImage) {
    workflow["5"] = { class_type: "EmptyLatentImage", inputs: { width, height, batch_size: resolveBatchSize(request) } };
    return finalizeSdxlWorkflow(workflow, samplerInputs, request);
  }

  workflow["4"] = { class_type: "LoadImage", inputs: { image: request.inputImage } };

  if (request.mode === "inpaint" && request.maskImage) {
    workflow["9"] = { class_type: "LoadImage", inputs: { image: request.maskImage } };
    workflow["11"] = { class_type: "ImageToMask", inputs: { image: ["9", 0], channel: "red" } };
    if (usesOriginalLatentContent(request)) {
      workflow["5"] = { class_type: "VAEEncode", inputs: { pixels: ["4", 0], vae: ["1", 2] } };
      workflow["12"] = { class_type: "GrowMask", inputs: { mask: ["11", 0], expand: resolveGrowMaskBy(request), tapered_corners: true } };
      workflow["13"] = { class_type: "SetLatentNoiseMask", inputs: { samples: ["5", 0], mask: ["12", 0] } };
      samplerInputs.latent_image = ["13", 0];
    } else {
      workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["4", 0], vae: ["1", 2], mask: ["11", 0], grow_mask_by: resolveGrowMaskBy(request) } };
    }
    return finalizeSdxlWorkflow(workflow, samplerInputs, request);
  }

  if (request.mode === "outpaint") {
    workflow["10"] = { class_type: "ImagePadForOutpaint", inputs: { image: ["4", 0], left: Math.round(request.outpaint?.left ?? 0), top: Math.round(request.outpaint?.top ?? 0), right: Math.round(request.outpaint?.right ?? 0), bottom: Math.round(request.outpaint?.bottom ?? 0), feathering: Math.round(request.outpaint?.feathering ?? 0) } };
    workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["10", 0], vae: ["1", 2], mask: ["10", 1], grow_mask_by: resolveGrowMaskBy(request) } };
    return finalizeSdxlWorkflow(workflow, samplerInputs, request);
  }

  workflow["5"] = { class_type: "VAEEncode", inputs: { pixels: ["4", 0], vae: ["1", 2] } };
  return finalizeSdxlWorkflow(workflow, samplerInputs, request);
}

function finalizeSdxlWorkflow(workflow: Workflow, samplerInputs: Record<string, unknown>, request: ComfyGenerateRequest): Workflow {
  const batchSize = resolveBatchSize(request);
  if (batchSize > 1 && workflow["5"]?.class_type !== "EmptyLatentImage") {
    const latent = samplerInputs.latent_image;
    workflow["19"] = { class_type: "RepeatLatentBatch", inputs: { samples: latent, amount: batchSize } };
    samplerInputs.latent_image = ["19", 0];
  }

  const control = request.inpaint?.structureControl ?? "none";
  if (request.mode !== "inpaint" || control === "none" || !request.inpaint?.controlModel || request.inpaint.controlModel === "auto") return addRefinementPass(workflow, samplerInputs, request);
  workflow["20"] = { class_type: "ControlNetLoader", inputs: { control_net_name: request.inpaint.controlModel } };
  if (control === "canny") {
    workflow["21"] = { class_type: "Canny", inputs: { image: ["4", 0], low_threshold: 100, high_threshold: 200 } };
  } else if (control === "depth") {
    workflow["21"] = { class_type: "MiDaS-DepthMapPreprocessor", inputs: { image: ["4", 0], a: 6.283, bg_threshold: 0.1, resolution: Math.max(request.width ?? 512, request.height ?? 512) } };
  } else {
    workflow["21"] = { class_type: "OpenposePreprocessor", inputs: { image: ["4", 0], detect_hand: "enable", detect_body: "enable", detect_face: "enable", scale_stick_for_xinsr_cn: "enable", resolution: Math.max(request.width ?? 512, request.height ?? 512) } };
  }
  workflow["22"] = {
    class_type: "ControlNetApplyAdvanced",
    inputs: {
      positive: ["2", 0],
      negative: ["3", 0],
      control_net: ["20", 0],
      image: ["21", 0],
      strength: Math.max(0, Math.min(1, request.inpaint.controlStrength ?? 0.55)),
      start_percent: 0,
      end_percent: 0.85,
      vae: ["1", 2],
    },
  };
  samplerInputs.positive = ["22", 0];
  samplerInputs.negative = ["22", 1];
  return addRefinementPass(workflow, samplerInputs, request);
}

function addRefinementPass(workflow: Workflow, samplerInputs: Record<string, unknown>, request: ComfyGenerateRequest): Workflow {
  if (!request.refinePass) return workflow;
  workflow["30"] = {
    class_type: "KSampler",
    inputs: {
      ...samplerInputs,
      seed: Math.round((request.seed ?? 0) + 1),
      steps: Math.max(6, Math.round((request.steps ?? 30) / 3)),
      denoise: Math.max(0, Math.min(1, (request.refineStrength ?? 20) / 100)),
      latent_image: ["6", 0],
    },
  };
  if (workflow["7"]) workflow["7"].inputs.samples = ["30", 0];
  return workflow;
}

function resolveBatchSize(request: ComfyGenerateRequest) {
  return Math.round(Math.max(1, Math.min(8, request.batchSize ?? 1)));
}

async function fetchObjectInfo(): Promise<ComfyObjectInfo> {
  const response = await fetch(`${comfyBaseUrl}/object_info`);
  if (!response.ok) throw new Error(`ComfyUI option lookup failed: ${response.status}`);
  return response.json() as Promise<ComfyObjectInfo>;
}

async function resolveStructureControl(request: ComfyGenerateRequest): Promise<ComfyGenerateRequest> {
  const control = request.inpaint?.structureControl ?? "none";
  if (request.mode !== "inpaint" || control === "none") return request;
  const info = await fetchObjectInfo();
  const requiredPreprocessor = control === "canny" ? info.Canny : control === "depth" ? info["MiDaS-DepthMapPreprocessor"] : info.OpenposePreprocessor;
  if (!requiredPreprocessor || !info.ControlNetLoader || !info.ControlNetApplyAdvanced) {
    throw new Error(`${control === "canny" ? "Canny" : control === "depth" ? "Depth" : "Pose"} structural control is not installed in ComfyUI.`);
  }
  const models = info.ControlNetLoader.input?.required?.control_net_name?.[0] ?? [];
  const requested = request.inpaint?.controlModel;
  const model = requested && requested !== "auto" ? requested : models.find((candidate) => candidate.toLowerCase().includes(control));
  if (!model || !models.includes(model)) throw new Error(`Install or select a ${control} ControlNet model before using structural control.`);
  return { ...request, inpaint: { ...request.inpaint, controlModel: model } };
}

export function buildZImageWorkflow(request: ComfyGenerateRequest): Workflow {
  return buildSeparatedTextToImageWorkflow(request, {
    architecture: "z-image",
    filenamePrefix: "image-studio-z-image",
    model: defaultModels["z-image"],
    textEncoder: "qwen_3_4b.safetensors",
    vae: "ae.safetensors",
    clipType: "lumina2",
    latentNode: "EmptySD3LatentImage",
    modelSamplingAuraFlow: true,
    negativeMode: "prompt",
    steps: 30,
    cfg: 4,
    sampler: "res_multistep",
    scheduler: "simple",
  });
}

export function buildZImageTurboWorkflow(request: ComfyGenerateRequest): Workflow {
  return buildSeparatedTextToImageWorkflow(request, {
    architecture: "z-image-turbo",
    filenamePrefix: "image-studio-z-image-turbo",
    model: defaultModels["z-image-turbo"],
    textEncoder: "qwen_3_4b.safetensors",
    vae: "ae.safetensors",
    clipType: "lumina2",
    latentNode: "EmptySD3LatentImage",
    modelSamplingAuraFlow: true,
    negativeMode: "zero",
    steps: 8,
    cfg: 1,
    sampler: "res_multistep",
    scheduler: "simple",
  });
}

export function buildAnimaWorkflow(request: ComfyGenerateRequest): Workflow {
  return buildSeparatedTextToImageWorkflow(request, {
    architecture: "anima",
    filenamePrefix: "image-studio-anima",
    model: defaultModels.anima,
    textEncoder: "qwen_3_06b_base.safetensors",
    vae: "qwen_image_vae.safetensors",
    clipType: "stable_diffusion",
    latentNode: "EmptyLatentImage",
    modelSamplingAuraFlow: false,
    negativeMode: "prompt",
    steps: 30,
    cfg: 4,
    sampler: "er_sde",
    scheduler: "simple",
  });
}

function buildSeparatedTextToImageWorkflow(request: ComfyGenerateRequest, config: {
  architecture: GenerateArchitecture;
  filenamePrefix: string;
  model: string;
  textEncoder: string;
  vae: string;
  clipType: string;
  latentNode: "EmptyLatentImage" | "EmptySD3LatentImage";
  modelSamplingAuraFlow: boolean;
  negativeMode: "prompt" | "zero";
  steps: number;
  cfg: number;
  sampler: string;
  scheduler: string;
}): Workflow {
  if (request.mode !== "text-to-image") throw new Error(`${architectureLabel(config.architecture)} currently supports text-to-image only`);
  const options = resolveSamplerOptions(request, config);
  const modelOutput: [string, number] = config.modelSamplingAuraFlow ? ["7", 0] : ["1", 0];
  const workflow: Workflow = {
    "1": { class_type: "UNETLoader", inputs: { unet_name: request.model && request.model !== "auto" ? request.model : config.model, weight_dtype: "default" } },
    "2": { class_type: "CLIPLoader", inputs: { clip_name: resolveSupportModelName(request.textEncoder, config.textEncoder), type: config.clipType, device: "default" } },
    "3": { class_type: "VAELoader", inputs: { vae_name: resolveSupportModelName(request.vae, config.vae) } },
    "4": { class_type: "CLIPTextEncode", inputs: { text: request.prompt, clip: ["2", 0] } },
    "6": { class_type: config.latentNode, inputs: { width: options.width, height: options.height, batch_size: 1 } },
    "8": { class_type: "KSampler", inputs: { seed: options.seed, steps: options.steps, cfg: options.cfg, sampler_name: options.sampler, scheduler: options.scheduler, denoise: 1, model: modelOutput, positive: ["4", 0], negative: ["5", 0], latent_image: ["6", 0] } },
    "9": { class_type: "VAEDecode", inputs: { samples: ["8", 0], vae: ["3", 0] } },
    "10": { class_type: "SaveImage", inputs: { filename_prefix: config.filenamePrefix, images: ["9", 0] } },
  };

  if (config.modelSamplingAuraFlow) workflow["7"] = { class_type: "ModelSamplingAuraFlow", inputs: { model: ["1", 0], shift: 3 } };
  workflow["5"] = config.negativeMode === "zero"
    ? { class_type: "ConditioningZeroOut", inputs: { conditioning: ["4", 0] } }
    : { class_type: "CLIPTextEncode", inputs: { text: request.negativePrompt ?? "", clip: ["2", 0] } };

  return workflow;
}

function resolveSamplerOptions(request: ComfyGenerateRequest, defaults: { steps: number; cfg: number; sampler: string; scheduler: string }) {
  if (!request.width || !request.height) throw new Error("Generation width and height are required");
  return {
    width: Math.max(64, Math.round(request.width)),
    height: Math.max(64, Math.round(request.height)),
    seed: request.seed === undefined || request.seed < 0 ? Math.floor(Math.random() * 2 ** 32) : Math.round(request.seed),
    steps: Math.max(1, Math.round(request.steps ?? defaults.steps)),
    cfg: Math.max(0, request.cfg ?? defaults.cfg),
    sampler: request.sampler ?? defaults.sampler,
    scheduler: request.scheduler ?? defaults.scheduler,
  };
}

function resolveSupportModelName(value: string | undefined, fallback: string) {
  return value && value !== "auto" ? value : fallback;
}

function resolveGrowMaskBy(request: ComfyGenerateRequest): number {
  const value = request.inpaint?.growMaskBy ?? 6;
  if (!Number.isFinite(value)) return 6;
  return Math.round(Math.max(0, Math.min(256, value)));
}

function usesOriginalLatentContent(request: ComfyGenerateRequest): boolean {
  return request.inpaint?.maskedContent === "original" || request.inpaint?.maskedContent === "originalColor" || request.inpaint?.maskedContent === "edges";
}

async function defaultModelForArchitecture(architecture: GenerateArchitecture, mode: GenerateMode) {
  if (architecture !== "sdxl") return defaultModels[architecture];
  const models = await listCheckpointModels();
  if (mode === "inpaint") return models.find((model) => /inpaint|fill/i.test(model)) ?? models[0] ?? defaultModels.sdxl;
  return models[0] ?? defaultModels.sdxl;
}

function normalizeArchitecture(architecture: ComfyGenerateRequest["architecture"]): GenerateArchitecture {
  if (architecture === "z-image" || architecture === "z-image-turbo" || architecture === "anima") return architecture;
  return "sdxl";
}

function architectureLabel(architecture: GenerateArchitecture): string {
  switch (architecture) {
    case "z-image":
      return "Z-Image";
    case "z-image-turbo":
      return "Z-Image Turbo";
    case "anima":
      return "Anima";
    case "sdxl":
    default:
      return "SDXL";
  }
}

function modelsForArchitecture(models: string[], defaultModel: string, matchers: RegExp[], options: { exclude?: RegExp[] } = {}) {
  return unique([defaultModel, ...models.filter((model) => {
    if (model === defaultModel) return true;
    if (options.exclude?.some((matcher) => matcher.test(model))) return false;
    if (matchers.length === 0) return true;
    return matchers.some((matcher) => matcher.test(model));
  })]);
}

function unique<T>(values: T[]) {
  return Array.from(new Set(values));
}

function parsePositiveInteger(value: string | undefined, fallback: number) {
  if (!value) return fallback;
  const parsed = Number(value);
  return Number.isFinite(parsed) && parsed > 0 ? Math.round(parsed) : fallback;
}

function nodeErrorsMessage(nodeErrors: unknown): string | undefined {
  if (!nodeErrors) return undefined;
  if (Array.isArray(nodeErrors) && nodeErrors.length === 0) return undefined;
  if (typeof nodeErrors === "object" && Object.keys(nodeErrors).length === 0) return undefined;

  if (typeof nodeErrors === "string") return nodeErrors;
  try {
    return JSON.stringify(nodeErrors);
  } catch {
    return "Unknown node validation error";
  }
}

function historyErrorMessage(history: unknown): string | undefined {
  const status = (history as { status?: { status_str?: string; completed?: boolean; messages?: unknown[] } }).status;
  if (!status) return undefined;
  if (status.status_str && status.status_str !== "success") return statusMessage(status);
  if (status.completed === false) return statusMessage(status);
  return undefined;
}

function statusMessage(status: { status_str?: string; messages?: unknown[] }) {
  const message = status.messages?.map(formatHistoryMessage).filter(Boolean).join("; ");
  return message || status.status_str || "Unknown execution error";
}

function formatHistoryMessage(message: unknown): string | undefined {
  if (!Array.isArray(message)) return undefined;
  const eventName = typeof message[0] === "string" ? message[0] : undefined;
  const payload = message[1];
  if (payload && typeof payload === "object") {
    const detail = payload as { exception_message?: string; node_type?: string; node_id?: string | number };
    if (detail.exception_message) {
      const node = detail.node_type ? ` in ${detail.node_type}${detail.node_id !== undefined ? ` ${detail.node_id}` : ""}` : "";
      return `${eventName ?? "error"}${node}: ${detail.exception_message}`;
    }
  }
  return eventName;
}