#!/usr/bin/env node

import fs from 'fs';
import path from 'path';
import crypto from 'crypto';

const DEFAULT_SERVER = process.env.COMFYUI_URL || 'http://192.168.1.202:8188';
const DEFAULT_OUT_ROOT = '/home/node/.openclaw/workspace/tmp/comfyui-workflow-tests';
const DEFAULT_POLL_MS = 3000;
const DEFAULT_TIMEOUT_SEC = 900;

function printUsage() {
  console.log(`Usage:
  node scripts/submit-flux2klein-edit-test.mjs --workflow <path> --image <path> [options]

Required:
  --workflow <path>         Flux2Klein UI workflow JSON with template subgraphs
  --image <path>            Primary reference image

Selection:
  --target single|multi     Select the 1-ref or 2-ref edit branch. Default: multi
  --node-id <id>            Override target branch by top-level node id
  --image2 <path>           Second reference image for multi mode

Overrides:
  --prompt <text>           Override the edit prompt
  --negative <text>         Override the negative prompt when the branch exposes one
  --seed <int>              Override noise seed
  --steps <int>             Override Flux2Scheduler steps
  --cfg <number>            Override CFG
  --sampler <name>          Override sampler
  --width <int>             Override output latent width
  --height <int>            Override output latent height
  --prefix <text>           SaveImage filename prefix

Execution:
  --server <url>            ComfyUI base URL. Default: ${DEFAULT_SERVER}
  --out-dir <path>          Output directory. Default: ${DEFAULT_OUT_ROOT}/<timestamp-slug>
  --poll-ms <int>           Poll interval in ms. Default: ${DEFAULT_POLL_MS}
  --timeout-sec <int>       Timeout in seconds. Default: ${DEFAULT_TIMEOUT_SEC}
  --no-wait                 Submit and exit after prompt_id is returned
  --dry-run                 Print the resolved prompt graph and exit
`);
}

function parseArgs(argv) {
  const options = {
    server: DEFAULT_SERVER,
    workflow: '',
    image: '',
    image2: '',
    target: 'multi',
    nodeId: '',
    prompt: '',
    negative: '',
    seed: null,
    steps: null,
    cfg: null,
    sampler: '',
    width: null,
    height: null,
    prefix: '',
    outDir: '',
    pollMs: DEFAULT_POLL_MS,
    timeoutSec: DEFAULT_TIMEOUT_SEC,
    dryRun: false,
    noWait: false,
  };

  for (let i = 0; i < argv.length; i += 1) {
    const arg = argv[i];
    if (arg === '--help' || arg === '-h') {
      printUsage();
      process.exit(0);
    }
    if (arg === '--dry-run') {
      options.dryRun = true;
      continue;
    }
    if (arg === '--no-wait') {
      options.noWait = true;
      continue;
    }
    const next = argv[i + 1];
    const take = () => {
      if (next === undefined) throw new Error(`Missing value for ${arg}`);
      i += 1;
      return next;
    };
    switch (arg) {
      case '--server': options.server = take(); break;
      case '--workflow': options.workflow = take(); break;
      case '--image': options.image = take(); break;
      case '--image2': options.image2 = take(); break;
      case '--target': options.target = take(); break;
      case '--node-id': options.nodeId = take(); break;
      case '--prompt': options.prompt = take(); break;
      case '--negative': options.negative = take(); break;
      case '--seed': options.seed = Number(take()); break;
      case '--steps': options.steps = Number(take()); break;
      case '--cfg': options.cfg = Number(take()); break;
      case '--sampler': options.sampler = take(); break;
      case '--width': options.width = Number(take()); break;
      case '--height': options.height = Number(take()); break;
      case '--prefix': options.prefix = take(); break;
      case '--out-dir': options.outDir = take(); break;
      case '--poll-ms': options.pollMs = Number(take()); break;
      case '--timeout-sec': options.timeoutSec = Number(take()); break;
      default:
        throw new Error(`Unknown argument: ${arg}`);
    }
  }

  if (!options.workflow) throw new Error('--workflow is required');
  if (!options.image) throw new Error('--image is required');
  return options;
}

function makeSlug(value) {
  return String(value || 'run')
    .toLowerCase()
    .replace(/[^a-z0-9]+/g, '-')
    .replace(/^-+|-+$/g, '')
    .slice(0, 80) || 'run';
}

function loadWorkflow(filePath) {
  const parsed = JSON.parse(fs.readFileSync(filePath, 'utf8'));
  if (!parsed || !Array.isArray(parsed.nodes) || !parsed.definitions?.subgraphs) {
    throw new Error(`Workflow is not a Flux2Klein UI workflow with subgraphs: ${filePath}`);
  }
  return parsed;
}

function normalizeLink(link) {
  if (Array.isArray(link)) {
    const [id, origin_id, origin_slot, target_id, target_slot, type] = link;
    return { id, origin_id, origin_slot, target_id, target_slot, type };
  }
  return link;
}

function indexLinks(links = []) {
  const byId = new Map();
  for (const raw of links) {
    const link = normalizeLink(raw);
    byId.set(link.id, link);
  }
  return { byId };
}

function pickTopLevelNode(workflow, options) {
  if (options.nodeId) {
    const direct = workflow.nodes.find((node) => String(node.id) === String(options.nodeId));
    if (!direct) throw new Error(`Top-level node ${options.nodeId} was not found`);
    return direct;
  }

  const candidateNodes = workflow.nodes.filter((node) => workflow.definitions.subgraphs.some((sg) => sg.id === node.type));
  const wantMulti = String(options.target || 'multi').toLowerCase() !== 'single';
  const picked = candidateNodes.find((node) => {
    const imageInputs = (node.inputs || []).filter((input) => String(input.type || '').toUpperCase() === 'IMAGE');
    return wantMulti ? imageInputs.length >= 2 : imageInputs.length === 1;
  });
  if (!picked) throw new Error(`Could not find a ${wantMulti ? 'multi' : 'single'} reference Flux2Klein branch in the workflow`);
  return picked;
}

function getSubgraph(workflow, topNode) {
  const subgraph = workflow.definitions.subgraphs.find((item) => item.id === topNode.type);
  if (!subgraph) throw new Error(`Subgraph definition ${topNode.type} was not found`);
  return subgraph;
}

function mapTopLevelInputs(topNode, subgraph, options) {
  const widgetValues = Array.isArray(topNode.widgets_values) ? [...topNode.widgets_values] : [];
  const mapped = {};

  for (let i = 0; i < subgraph.inputs.length; i += 1) {
    const input = subgraph.inputs[i];
    mapped[input.name] = widgetValues[i];
  }

  if (options.prompt) mapped.text = options.prompt;
  if (options.seed !== null) mapped.noise_seed = Math.trunc(options.seed);
  if (options.image) mapped.image = options.image;
  if (options.image2) mapped.image_1 = options.image2;
  return mapped;
}

async function fetchObjectInfo(server, nodeName) {
  const response = await fetch(`${server.replace(/\/$/, '')}/object_info/${encodeURIComponent(nodeName)}`);
  if (!response.ok) throw new Error(`Failed to fetch object info for ${nodeName} (${response.status})`);
  const payload = await response.json();
  return payload?.[nodeName] || null;
}

function visibleInputNames(objectInfo) {
  const order = objectInfo?.input_order || {};
  return [
    ...(order.required || []),
    ...(order.optional || []),
  ];
}

function getInputSpec(objectInfo, inputName) {
  return objectInfo?.input?.required?.[inputName]
    || objectInfo?.input?.optional?.[inputName]
    || null;
}

function isConnectionType(specType) {
  return new Set(['MODEL', 'IMAGE', 'LATENT', 'VAE', 'CONDITIONING', 'SIGMAS', 'GUIDER', 'SAMPLER', 'NOISE', 'CLIP']).has(specType);
}

function isWidgetBacked(entry, inputName, orderedNames, objectInfo) {
  if (entry?.widget) return true;
  const spec = getInputSpec(objectInfo, inputName);
  if (!spec) return false;
  const specType = Array.isArray(spec) ? spec[0] : spec;
  if (Array.isArray(specType)) return true;
  return orderedNames.includes(inputName) && !isConnectionType(String(specType || ''));
}

function normalizeModelToken(value) {
  return String(value || '')
    .replace(/\\/g, '/')
    .split('/')
    .pop()
    .toLowerCase()
    .replace(/[^a-z0-9]+/g, '');
}

function pickCompatibleModel(currentValue, available) {
  if (available.includes(currentValue)) return currentValue;
  const lower = available.find((item) => String(item).toLowerCase() === String(currentValue).toLowerCase());
  if (lower) return lower;
  const currentBase = path.basename(String(currentValue || '').replace(/\\/g, '/'));
  const byBase = available.find((item) => path.basename(String(item).replace(/\\/g, '/')).toLowerCase() === currentBase.toLowerCase());
  if (byBase) return byBase;
  const aliases = new Map([
    ['flux-2-klein-9b-fp8.safetensors', 'flux-2-klein-9b-fp8mixed.safetensors'],
    ['flux-2-klein-4b-fp8.safetensors', 'flux-2-klein-4b-fp8mixed.safetensors'],
    ['qwen_3_4b.safetensors', 'Qwen\\qwen_3_4b.safetensors'],
  ]);
  const aliased = aliases.get(String(currentValue || ''));
  if (aliased && available.includes(aliased)) return aliased;
  const normalized = normalizeModelToken(currentValue);
  return available.find((item) => normalizeModelToken(item) === normalized) || null;
}

async function remapLoaderValue(server, classType, inputName, currentValue) {
  const info = await fetchObjectInfo(server, classType);
  const available = info?.input?.required?.[inputName]?.[0];
  if (!Array.isArray(available) || available.includes(currentValue)) return currentValue;
  const replacement = pickCompatibleModel(currentValue, available);
  if (!replacement) {
    throw new Error(`${classType}.${inputName} value '${currentValue}' is not available on the live server`);
  }
  return replacement;
}

async function uploadImageToComfyUi(server, filePath) {
  const buffer = fs.readFileSync(filePath);
  const fileName = path.basename(filePath);
  const form = new FormData();
  form.append('image', new Blob([buffer], { type: guessMimeType(filePath) }), fileName);
  form.append('type', 'input');
  form.append('subfolder', '');
  form.append('overwrite', 'true');
  const response = await fetch(`${server.replace(/\/$/, '')}/api/upload/image`, {
    method: 'POST',
    body: form,
  });
  const text = await response.text();
  const payload = text ? JSON.parse(text) : {};
  if (!response.ok) {
    throw new Error(payload?.error || payload?.message || `Upload failed (${response.status})`);
  }
  return String(payload?.name || payload?.filename || payload?.file || fileName);
}

function guessMimeType(filePath) {
  const ext = path.extname(filePath).toLowerCase();
  if (ext === '.png') return 'image/png';
  if (ext === '.webp') return 'image/webp';
  return 'image/jpeg';
}

function buildSubgraphPromptGraph(subgraph, externalValues, objectInfoByType, overrides = {}, extraNodes = {}) {
  const { byId } = indexLinks(subgraph.links || []);
  const prompt = { ...extraNodes };
  const outputRefs = new Map();

  for (const node of subgraph.nodes) {
    const nodeKey = String(node.id);
    const objectInfo = objectInfoByType.get(node.type);
    if (!objectInfo) throw new Error(`Missing object info for ${node.type}`);
    const orderedNames = visibleInputNames(objectInfo);
    const inputs = {};
    let widgetIndex = 0;

    for (const inputName of orderedNames) {
      const entry = (node.inputs || []).find((item) => item.name === inputName);
      const link = entry?.link != null ? byId.get(entry.link) : null;
      const shouldConsumeWidget = isWidgetBacked(entry, inputName, orderedNames, objectInfo);
      const widgetValue = shouldConsumeWidget && widgetIndex < (node.widgets_values || []).length
        ? node.widgets_values[widgetIndex++]
        : undefined;

      if (link) {
        if (Number(link.origin_id) === -10) {
          const external = subgraph.inputs[Number(link.origin_slot)];
          inputs[inputName] = externalValues[external.name];
        } else {
          inputs[inputName] = [String(link.origin_id), Number(link.origin_slot)];
        }
      } else if (widgetValue !== undefined) {
        inputs[inputName] = widgetValue;
      }
    }

    for (const entry of node.inputs || []) {
      if (Object.prototype.hasOwnProperty.call(inputs, entry.name)) continue;
      if (entry.link == null) continue;
      const link = byId.get(entry.link);
      if (!link) continue;
      if (Number(link.origin_id) === -10) {
        const external = subgraph.inputs[Number(link.origin_slot)];
        inputs[entry.name] = externalValues[external.name];
      } else {
        inputs[entry.name] = [String(link.origin_id), Number(link.origin_slot)];
      }
    }

    prompt[nodeKey] = {
      class_type: node.type,
      inputs,
      _meta: { title: node.title || node.type },
    };
    if (node.type === 'RandomNoise' && prompt[nodeKey].inputs.noise_seed === undefined && Array.isArray(node.widgets_values) && node.widgets_values.length) {
      prompt[nodeKey].inputs.noise_seed = node.widgets_values[0];
    }
  }

  for (const link of byId.values()) {
    if (Number(link.target_id) !== -20) continue;
    outputRefs.set(Number(link.target_slot), [String(link.origin_id), Number(link.origin_slot)]);
  }

  for (const [nodeId, promptNode] of Object.entries(prompt)) {
    if (promptNode.class_type === 'KSamplerSelect' && overrides.sampler) {
      promptNode.inputs.sampler_name = overrides.sampler;
    }
    if (promptNode.class_type === 'Flux2Scheduler') {
      if (overrides.steps !== null) promptNode.inputs.steps = Math.trunc(overrides.steps);
      if (overrides.width !== null) promptNode.inputs.width = Math.trunc(overrides.width);
      if (overrides.height !== null) promptNode.inputs.height = Math.trunc(overrides.height);
    }
    if (promptNode.class_type === 'EmptyFlux2LatentImage') {
      if (overrides.width !== null) promptNode.inputs.width = Math.trunc(overrides.width);
      if (overrides.height !== null) promptNode.inputs.height = Math.trunc(overrides.height);
    }
    if (promptNode.class_type === 'CFGGuider' && overrides.cfg !== null) {
      promptNode.inputs.cfg = overrides.cfg;
    }
    if (promptNode.class_type === 'RandomNoise' && overrides.seed !== null) {
      promptNode.inputs.noise_seed = Math.trunc(overrides.seed);
    }
    if (promptNode.class_type === 'CLIPTextEncode' && typeof promptNode.inputs.text === 'string') {
      if (overrides.prompt && promptNode.inputs.text && !overrides.negativeApplied) {
        promptNode.inputs.text = overrides.prompt;
        overrides.negativeApplied = true;
      } else if (!promptNode.inputs.text && overrides.negative) {
        promptNode.inputs.text = overrides.negative;
      }
    }
  }

  const saveNodeId = String(Math.max(...Object.keys(prompt).map((value) => Number(value))) + 1000);
  const primaryOutput = outputRefs.get(0);
  if (!primaryOutput) throw new Error('The selected subgraph does not expose an IMAGE output');
  prompt[saveNodeId] = {
    class_type: 'SaveImage',
    inputs: {
      images: primaryOutput,
      filename_prefix: overrides.prefix || 'Flux2-Klein',
    },
    _meta: { title: 'Save Image' },
  };

  return prompt;
}

async function buildPromptFromTemplate(workflow, topNode, subgraph, options) {
  const externalValues = mapTopLevelInputs(topNode, subgraph, options);
  const uploadedPrimary = await uploadImageToComfyUi(options.server, path.resolve(options.image));
  const uploadedSecondary = externalValues.image_1
    ? await uploadImageToComfyUi(options.server, path.resolve(externalValues.image_1))
    : null;

  const loadImageInfo = await fetchObjectInfo(options.server, 'LoadImage');
  const objectInfoByType = new Map([['LoadImage', loadImageInfo]]);
  const extraNodes = {};
  const imageInputDefs = subgraph.inputs.filter((input) => String(input.type || '').toUpperCase() === 'IMAGE');
  let nextExternalNodeId = 900000;
  for (const inputDef of imageInputDefs) {
    const uploadedName = inputDef.name === 'image_1' ? uploadedSecondary : uploadedPrimary;
    if (!uploadedName) continue;
    const loadNodeId = String(nextExternalNodeId++);
    externalValues[inputDef.name] = [loadNodeId, 0];
    extraNodes[loadNodeId] = {
      class_type: 'LoadImage',
      inputs: {
        image: uploadedName,
      },
      _meta: { title: `Load Image (${inputDef.name})` },
    };
  }

  const typeNames = new Set(subgraph.nodes.map((node) => node.type));
  typeNames.add('SaveImage');
  for (const typeName of typeNames) {
    if (objectInfoByType.has(typeName)) continue;
    objectInfoByType.set(typeName, await fetchObjectInfo(options.server, typeName));
  }

  const prompt = buildSubgraphPromptGraph(subgraph, externalValues, objectInfoByType, {
    prompt: options.prompt,
    negative: options.negative,
    seed: options.seed,
    steps: options.steps,
    cfg: options.cfg,
    sampler: options.sampler,
    width: options.width,
    height: options.height,
    prefix: options.prefix,
    negativeApplied: false,
  }, extraNodes);

  for (const node of Object.values(prompt)) {
    if (node.class_type === 'UNETLoader' && node.inputs.unet_name) {
      node.inputs.unet_name = await remapLoaderValue(options.server, 'UNETLoader', 'unet_name', node.inputs.unet_name);
    }
    if (node.class_type === 'CLIPLoader' && node.inputs.clip_name) {
      node.inputs.clip_name = await remapLoaderValue(options.server, 'CLIPLoader', 'clip_name', node.inputs.clip_name);
    }
    if (node.class_type === 'VAELoader' && node.inputs.vae_name) {
      node.inputs.vae_name = await remapLoaderValue(options.server, 'VAELoader', 'vae_name', node.inputs.vae_name);
    }
  }

  return {
    prompt,
    uploadedImages: {
      image: uploadedPrimary,
      image2: uploadedSecondary,
    },
  };
}

async function postPrompt(server, promptGraph) {
  const response = await fetch(`${server.replace(/\/$/, '')}/prompt`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({
      client_id: crypto.randomUUID(),
      prompt: promptGraph,
    }),
  });
  if (!response.ok) {
    const text = await response.text();
    throw new Error(`ComfyUI prompt submission failed (${response.status}): ${text}`);
  }
  return response.json();
}

function normalizeHistoryRecord(raw, promptId) {
  if (!raw || typeof raw !== 'object') return null;
  if (raw[promptId]) return raw[promptId];
  if (raw.prompt_id || raw.status || raw.outputs) return raw;
  const firstKey = Object.keys(raw)[0];
  return firstKey ? raw[firstKey] : null;
}

async function waitForHistory(server, promptId, pollMs, timeoutSec) {
  const deadline = Date.now() + (timeoutSec * 1000);
  while (Date.now() < deadline) {
    const response = await fetch(`${server.replace(/\/$/, '')}/history/${encodeURIComponent(promptId)}`);
    if (response.ok) {
      const payload = await response.json();
      const record = normalizeHistoryRecord(payload, promptId);
      const completed = Boolean(record?.status?.completed);
      const failed = String(record?.status?.status_str || '').toLowerCase() === 'error';
      if (completed || failed) return record;
    }
    await new Promise((resolve) => setTimeout(resolve, pollMs));
  }
  throw new Error(`Timed out waiting for ComfyUI history for prompt ${promptId}`);
}

function collectImageEntries(historyRecord) {
  const images = [];
  const outputs = historyRecord?.outputs || {};
  for (const [nodeId, nodeOutput] of Object.entries(outputs)) {
    if (!nodeOutput || typeof nodeOutput !== 'object') continue;
    const entries = Array.isArray(nodeOutput.images) ? nodeOutput.images : [];
    for (const entry of entries) {
      images.push({ nodeId, ...entry });
    }
  }
  return images;
}

async function downloadOutputImages(server, imageEntries, outDir) {
  const saved = [];
  fs.mkdirSync(outDir, { recursive: true });
  for (let i = 0; i < imageEntries.length; i += 1) {
    const image = imageEntries[i];
    const params = new URLSearchParams({
      filename: String(image.filename || ''),
      subfolder: String(image.subfolder || ''),
      type: String(image.type || 'output'),
    });
    const response = await fetch(`${server.replace(/\/$/, '')}/view?${params.toString()}`);
    if (!response.ok) throw new Error(`Failed to download image ${image.filename} (${response.status})`);
    const buffer = Buffer.from(await response.arrayBuffer());
    const ext = path.extname(String(image.filename || '')) || '.png';
    const localName = `${String(i + 1).padStart(2, '0')}-${path.basename(String(image.filename || `output${ext}`))}`;
    const localPath = path.join(outDir, localName);
    fs.writeFileSync(localPath, buffer);
    saved.push({ ...image, localPath });
  }
  return saved;
}

async function main() {
  const options = parseArgs(process.argv.slice(2));
  const workflowPath = path.resolve(options.workflow);
  const workflow = loadWorkflow(workflowPath);
  const topNode = pickTopLevelNode(workflow, options);
  const subgraph = getSubgraph(workflow, topNode);
  const { prompt, uploadedImages } = await buildPromptFromTemplate(workflow, topNode, subgraph, options);

  const runSlug = makeSlug(`${path.basename(workflowPath, path.extname(workflowPath))}-${options.target}`);
  const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
  const outDir = path.resolve(options.outDir || path.join(DEFAULT_OUT_ROOT, `${timestamp}-${runSlug}`));

  if (options.dryRun) {
    console.log(JSON.stringify({ prompt, uploadedImages, topNodeId: topNode.id, subgraphId: subgraph.id }, null, 2));
    return;
  }

  fs.mkdirSync(outDir, { recursive: true });
  fs.writeFileSync(path.join(outDir, 'resolved-prompt.json'), JSON.stringify(prompt, null, 2));

  const submission = await postPrompt(options.server, prompt);
  const promptId = String(submission?.prompt_id || submission?.promptId || '').trim();
  if (!promptId) throw new Error(`ComfyUI did not return a prompt_id: ${JSON.stringify(submission)}`);

  const runRecord = {
    workflowPath,
    server: options.server,
    promptId,
    submittedAt: new Date().toISOString(),
    topNodeId: topNode.id,
    subgraphId: subgraph.id,
    target: options.target,
    uploadedImages,
    overrides: {
      prompt: options.prompt || null,
      negative: options.negative || null,
      seed: options.seed,
      steps: options.steps,
      cfg: options.cfg,
      sampler: options.sampler || null,
      width: options.width,
      height: options.height,
      prefix: options.prefix || null,
    },
    outputDir: outDir,
  };

  if (options.noWait) {
    fs.writeFileSync(path.join(outDir, 'run-record.json'), JSON.stringify({ ...runRecord, status: 'submitted-no-wait' }, null, 2));
    console.log(JSON.stringify({
      ok: true,
      promptId,
      outputDir: outDir,
      status: 'submitted-no-wait',
    }, null, 2));
    return;
  }

  const history = await waitForHistory(options.server, promptId, options.pollMs, options.timeoutSec);
  fs.writeFileSync(path.join(outDir, 'history.json'), JSON.stringify(history, null, 2));
  const images = collectImageEntries(history);
  const savedImages = await downloadOutputImages(options.server, images, outDir);
  fs.writeFileSync(path.join(outDir, 'run-record.json'), JSON.stringify({
    ...runRecord,
    imageCount: savedImages.length,
    images: savedImages,
    status: history?.status || null,
  }, null, 2));

  console.log(JSON.stringify({
    ok: true,
    promptId,
    outputDir: outDir,
    imageCount: savedImages.length,
    images: savedImages.map((item) => item.localPath),
    status: history?.status?.status_str || 'unknown',
  }, null, 2));
}

main().catch((error) => {
  console.error(error?.stack || error?.message || String(error));
  process.exit(1);
});