#!/usr/bin/env node

import fs from 'fs';
import path from 'path';
import { execFileSync } from 'child_process';

const REPO_ROOT = '/home/node/.openclaw/workspace/work/comfyui-workflows';
const DEFAULT_SERVER = process.env.COMFYUI_URL || 'http://192.168.1.202:8188';
const DEFAULT_WORKFLOW = path.join(REPO_ROOT, 'models/flux2klein/image-edit-9b-distilled/workflow.json');
const DEFAULT_INPUT = '/home/node/.openclaw/workspace/tmp/flux2klein-inputs/2026-08-13-pinup-ref.jpg';
const DEFAULT_STATE_FILE = path.join(REPO_ROOT, 'benchmarks/flux2klein/overnight-state.json');
const DEFAULT_RUNS_ROOT = path.join(REPO_ROOT, 'benchmarks/flux2klein/runs');
const SUBMIT_SCRIPT = path.join(REPO_ROOT, 'scripts/submit-flux2klein-edit-test.mjs');

const DEFAULT_PROMPT = 'Restage the exact same woman as a clean full-body studio reference photo on a pure white seamless background. Face the camera directly in a neutral standing pose. Preserve identical facial features, platinum blonde hair, skin tone, gold corset dress, black thigh-high boots, proportions, and age. Even soft studio lighting, no props, no street, no buildings, no text.';

function printUsage() {
  console.log(`Usage:
  node scripts/run-flux2klein-overnight-smoke.mjs [options]

Options:
  --server <url>             ComfyUI base URL. Default: ${DEFAULT_SERVER}
  --workflow <path>          Workflow path. Default: ${DEFAULT_WORKFLOW}
  --image <path>             Input image path. Default: ${DEFAULT_INPUT}
  --prompt <text>            Prompt override
  --steps <int>              Default: 4
  --cfg <number>             Default: 1.0
  --width <int>              Default: 1536
  --height <int>             Default: 864
  --prefix <text>            Output filename prefix. Default: flux2klein-overnight-smoke
  --batch-label <text>       Batch label appended to the run folder name
  --state-file <path>        State file path. Default: ${DEFAULT_STATE_FILE}
  --optimized-path <path>    Relative repo path touched by the optimisation pass. Repeatable.
  --help                     Show this message
`);
}

function parseArgs(argv) {
  const options = {
    server: DEFAULT_SERVER,
    workflow: DEFAULT_WORKFLOW,
    image: DEFAULT_INPUT,
    prompt: DEFAULT_PROMPT,
    steps: 4,
    cfg: 1.0,
    width: 1536,
    height: 864,
    prefix: 'flux2klein-overnight-smoke',
    batchLabel: 'single-smoke',
    stateFile: DEFAULT_STATE_FILE,
    optimizedPaths: [],
  };

  for (let i = 0; i < argv.length; i += 1) {
    const arg = argv[i];
    if (arg === '--help' || arg === '-h') {
      printUsage();
      process.exit(0);
    }
    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 '--prompt': options.prompt = take(); break;
      case '--steps': options.steps = Number(take()); break;
      case '--cfg': options.cfg = Number(take()); break;
      case '--width': options.width = Number(take()); break;
      case '--height': options.height = Number(take()); break;
      case '--prefix': options.prefix = take(); break;
      case '--batch-label': options.batchLabel = take(); break;
      case '--state-file': options.stateFile = take(); break;
      case '--optimized-path': options.optimizedPaths.push(take()); break;
      default:
        throw new Error(`Unknown argument: ${arg}`);
    }
  }

  return options;
}

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

function repoRelative(filePath) {
  return path.relative(REPO_ROOT, filePath).replace(/\\/g, '/');
}

function ensureDir(filePath) {
  fs.mkdirSync(filePath, { recursive: true });
}

function readJson(filePath) {
  return JSON.parse(fs.readFileSync(filePath, 'utf8'));
}

function writeJson(filePath, data) {
  ensureDir(path.dirname(filePath));
  fs.writeFileSync(filePath, `${JSON.stringify(data, null, 2)}\n`);
}

async function fetchQueue(server) {
  const base = server.replace(/\/$/, '');
  const endpoints = [`${base}/api/queue`, `${base}/queue`];
  let lastError = null;
  for (const endpoint of endpoints) {
    try {
      const response = await fetch(endpoint);
      if (!response.ok) {
        lastError = new Error(`Queue check failed at ${endpoint} (${response.status})`);
        continue;
      }
      return await response.json();
    } catch (error) {
      lastError = error;
    }
  }
  throw lastError || new Error('Queue check failed');
}

function summarizeBusy(queueData) {
  const running = Array.isArray(queueData?.queue_running) ? queueData.queue_running.length : 0;
  const pending = Array.isArray(queueData?.queue_pending) ? queueData.queue_pending.length : 0;
  return {
    running,
    pending,
    busy: running > 0 || pending > 0,
  };
}

function writeState(options, patch) {
  const previous = fs.existsSync(options.stateFile) ? readJson(options.stateFile) : {};
  const next = {
    testsSet: false,
    serverBusy: false,
    reason: '',
    cases: [],
    reviewPending: false,
    verdict: null,
    ...previous,
    ...patch,
    date: new Date().toISOString(),
    optimizedPaths: options.optimizedPaths,
  };
  writeJson(options.stateFile, next);
  return next;
}

function copyInputImage(sourcePath, caseDir) {
  const ext = path.extname(sourcePath) || '.bin';
  const destination = path.join(caseDir, `input-1${ext}`);
  fs.copyFileSync(sourcePath, destination);
  return destination;
}

function runSubmitScript(options, caseDir) {
  const stdout = execFileSync('node', [
    SUBMIT_SCRIPT,
    '--workflow', options.workflow,
    '--image', options.image,
    '--target', 'single',
    '--prompt', options.prompt,
    '--steps', String(options.steps),
    '--cfg', String(options.cfg),
    '--width', String(options.width),
    '--height', String(options.height),
    '--prefix', options.prefix,
    '--out-dir', caseDir,
    '--server', options.server,
  ], {
    cwd: REPO_ROOT,
    encoding: 'utf8',
    stdio: ['ignore', 'pipe', 'pipe'],
  });
  return JSON.parse(stdout);
}

async function main() {
  const options = parseArgs(process.argv.slice(2));
  options.workflow = path.resolve(options.workflow);
  options.image = path.resolve(options.image);
  options.stateFile = path.resolve(options.stateFile);

  if (!fs.existsSync(options.workflow)) {
    const state = writeState(options, {
      reason: `Workflow not found: ${repoRelative(options.workflow)}`,
    });
    console.log(JSON.stringify({ ok: false, state }, null, 2));
    return;
  }

  if (!fs.existsSync(options.image)) {
    const state = writeState(options, {
      reason: `Input image not found: ${options.image}`,
    });
    console.log(JSON.stringify({ ok: false, state }, null, 2));
    return;
  }

  let queueData;
  try {
    queueData = await fetchQueue(options.server);
  } catch (error) {
    const state = writeState(options, {
      reason: `Could not verify queue state: ${error.message}`,
    });
    console.log(JSON.stringify({ ok: false, state }, null, 2));
    return;
  }

  const queue = summarizeBusy(queueData);
  if (queue.busy) {
    const state = writeState(options, {
      serverBusy: true,
      reason: `Skipped smoke test because ComfyUI was busy (running=${queue.running}, pending=${queue.pending}).`,
      cases: [],
      reviewPending: false,
    });
    console.log(JSON.stringify({ ok: true, skipped: true, queue, state }, null, 2));
    return;
  }

  const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
  const batchDir = path.join(DEFAULT_RUNS_ROOT, `${timestamp}-${makeSlug(options.batchLabel)}`);
  const caseDir = path.join(batchDir, 'cases', 'case01');
  ensureDir(caseDir);
  const copiedInput = copyInputImage(options.image, caseDir);

  const result = runSubmitScript(options, caseDir);
  const runRecordPath = path.join(caseDir, 'run-record.json');
  const promptPath = path.join(caseDir, 'resolved-prompt.json');
  const runRecord = readJson(runRecordPath);
  const outputEntries = Array.isArray(runRecord.images) && runRecord.images.length
    ? runRecord.images
    : Array.isArray(runRecord.files) && runRecord.files.length
      ? runRecord.files
      : [];
  const outputFile = outputEntries.length
    ? outputEntries[0].localPath
    : null;

  const manifest = {
    workflow: repoRelative(options.workflow),
    input: repoRelative(copiedInput),
    output: outputFile ? repoRelative(outputFile) : null,
    runRecord: repoRelative(runRecordPath),
    resolvedPrompt: repoRelative(promptPath),
    promptId: result.promptId || runRecord.promptId || null,
    status: result.status || runRecord?.status?.status_str || 'unknown',
    optimizedPaths: options.optimizedPaths,
    overrides: {
      steps: options.steps,
      cfg: options.cfg,
      width: options.width,
      height: options.height,
      prefix: options.prefix,
    },
  };
  writeJson(path.join(batchDir, 'manifest.json'), manifest);

  const state = writeState(options, {
    testsSet: true,
    serverBusy: false,
    reason: 'Overnight flux2klein smoke test submitted successfully.',
    cases: [manifest],
    reviewPending: true,
    verdict: null,
  });

  console.log(JSON.stringify({
    ok: true,
    queue,
    batchDir: repoRelative(batchDir),
    caseDir: repoRelative(caseDir),
    state,
  }, null, 2));
}

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