#!/usr/bin/env node
import path from 'node:path';
import {fileURLToPath} from 'node:url';
import {loadConfig, resolveStage, publicStage, samInfo} from '../lib/config.mjs';
import {requestJson} from '../providers/http.mjs';
import {runUnderstanding, runCompletionProbe} from '../stages/model-stage.mjs';
import {readImage} from '../lib/images.mjs';
import {runSegmentation} from '../stages/segmentation-stage.mjs';
import {runCompletion} from '../stages/completion-stage.mjs';

const help = `模型接入与结构理解（新增模式，不调用旧拆件流程）
  node pipeline.mjs --mode annotated-parts --reference image.png --out-dir output/run --stop-after understand
  node tools/model-check.mjs --stage completion --image local.png --prompt-file task.txt --out-dir output/edit
  node tools/model-check.mjs --stage segmentation --input-dir output/understand --out-dir output/segment --sam-dir F:/Study/SAM3
  node tools/model-check.mjs --stage completion --input-dir output/segment --out-dir output/complete --part head --prepare-only
  node tools/model-check.mjs --list-models
  node tools/model-check.mjs --dry-run

公共选项：
  --config FILE                  默认 configs/astra-qwen.json（prompt 相对配置文件）
  --env-file FILE                默认项目 .env，shell 环境变量优先
  --understanding-model ID       理解模型覆盖（使用原样 ID，不自动映射）
  --completion-model ID          图片编辑模型覆盖
  --understanding-protocol NAME  chat-completions / responses
  --stream true|false            覆盖理解阶段的流式设置
  --output-mode NAME             text-json / json-object / json-schema
  --resume                      输入/配置/代码和产物哈希相同时复用
  --parts-file FILE              离线读取已有 v1/v2 计划，做 Alpha 校验与种子生成，不请求模型
  --python FILE                  ai-spine Conda 环境解释器（默认使用已激活环境）
  --input-dir DIR                SAM 阶段读理解目录；补全阶段读 SAM 目录
  --part ID                      补全部件 ID，可重复；默认处理所有需补全部件
  --prepare-only                 仅准备补全画布、掩码和提示词，不请求模型
  --candidates-dir DIR           离线导入 ID.png，或旧补全目录的首张候选，不请求模型
  --candidate-size-policy NAME   strict（默认）/ uniform（同宽高比缩放后重新检查对齐）
  --completion-size INT          补全方形画布边长，默认 1024，不缩放原像素
  --extension-radius INT         补全区域距可见轮廓的最大像素距离，默认 32
  --sam-dir DIR                  本地 SAM 根目录，包含 sam3 与 sam3.1checkpoint
  --sam-code-dir DIR             SAM 源码仓库目录覆盖
  --sam-checkpoint FILE          本地权重路径覆盖
  --sam-confidence FLOAT         模型候选分数阈值，默认 0.2
  --sam-precision NAME           bf16（默认）或 fp32
  --image FILE                  图片编辑输入，可重复；多图需配置 capabilities
  --mask FILE                   透明处可编辑的 RGBA PNG；需配置 capabilities
  --dry-run                     检查配置与本地输入，不请求模型、不生成产物

--stop-after 支持 understand / segment。SAM 在临时 GPU 子进程运行，完成后退出。
部件补全使用 --stage completion --input-dir；可见像素恢复与几何检查已接入，Spine 导出尚未接入。`;

export function parseOptions(argv) {
  const options = {stage: 'understanding', imagePaths: [], partIds: []};
  const valueFlags = new Map([
    ['--config', 'configFile'], ['--env-file', 'envFile'], ['--reference', 'reference'], ['--out-dir', 'outDir'],
    ['--understanding-model', 'understandingModel'], ['--completion-model', 'completionModel'],
    ['--understanding-protocol', 'understandingProtocol'], ['--output-mode', 'outputMode'],
    ['--stage', 'stage'], ['--mask', 'maskPath'], ['--prompt-file', 'promptFile'], ['--stop-after', 'stopAfter'],
    ['--mode', 'mode'], ['--python', 'python'], ['--stream', 'streamRaw'], ['--parts-file', 'partsFile'],
    ['--input-dir', 'inputDir'], ['--sam-dir', 'samDir'], ['--sam-code-dir', 'samCodeDir'],
    ['--sam-checkpoint', 'samCheckpoint'], ['--sam-confidence', 'samConfidence'], ['--sam-precision', 'samPrecision'],
    ['--candidates-dir', 'candidatesDir'], ['--completion-size', 'completionSize'], ['--extension-radius', 'extensionRadius'],
    ['--candidate-size-policy', 'candidateSizePolicy'],
  ]);
  for (let i = 0; i < argv.length; i += 1) {
    const flag = argv[i];
    if (['--help', '-h'].includes(flag)) options.help = true;
    else if (flag === '--resume') options.resume = true;
    else if (flag === '--dry-run') options.dryRun = true;
    else if (flag === '--prepare-only') options.prepareOnly = true;
    else if (flag === '--list-models') options.listModels = true;
    else if (valueFlags.has(flag) || flag === '--image' || flag === '--part') {
      const value = argv[++i];
      if (!value || value.startsWith('--')) throw new Error(`${flag} 缺少参数值`);
      if (flag === '--image') options.imagePaths.push(value);
      else if (flag === '--part') options.partIds.push(value);
      else options[valueFlags.get(flag)] = value;
    } else throw new Error(`未知参数：${flag}，请用 --help 查看`);
  }
  return options;
}

export async function main(argv = process.argv.slice(2)) {
  const options = parseOptions(argv);
  if (options.help) { console.log(help); return; }
  if (!['understanding', 'completion', 'segmentation'].includes(options.stage)) throw new Error('--stage 仅支持 understanding / completion / segmentation');
  if (options.partsFile && (options.stage !== 'understanding' || options.listModels)) throw new Error('--parts-file 仅用于离线结构理解');
  if (options.streamRaw !== undefined && !['true', 'false'].includes(options.streamRaw)) throw new Error('--stream 仅支持 true 或 false');
  options.stream = options.streamRaw === undefined ? undefined : options.streamRaw === 'true';
  if (options.mode && options.mode !== 'annotated-parts') throw new Error('新入口仅支持 --mode annotated-parts');
  if (options.stopAfter && (!['understand', 'segment'].includes(options.stopAfter) || options.stage === 'completion' || (options.stage === 'segmentation' && options.stopAfter !== 'segment'))) {
    throw new Error('--stop-after 仅支持 understand / segment，且必须匹配阶段');
  }
  const loaded = loadConfig(options);
  const partCompletion = options.stage === 'completion' && options.inputDir;
  if (!partCompletion && (options.prepareOnly || options.candidatesDir || options.partIds.length || options.completionSize || options.extensionRadius || options.candidateSizePolicy)) {
    throw new Error('部件补全参数需要 --stage completion --input-dir');
  }
  if (partCompletion) {
    if (options.reference || options.imagePaths.length || options.maskPath || options.promptFile || options.listModels) throw new Error('部件补全从 --input-dir 生成画布与提示，不接受探针输入参数');
    const stage = resolveStage(loaded, 'completion', {model: options.completionModel}, {requireKey: !options.dryRun && !options.prepareOnly && !options.candidatesDir});
    const result = await runCompletion({...options, stage});
    if (!result.dryRun) console.log(`部件补全${result.reused ? '复用完成' : options.prepareOnly ? '准备完成（未请求模型）' : '处理完成，需复核'}：${path.join(result.out, 'report.html')}`);
    return result;
  }
  if (options.stage === 'segmentation' || (options.stopAfter === 'segment' && options.inputDir)) {
    if (options.reference || options.partsFile || options.imagePaths.length || options.maskPath || options.promptFile || options.listModels) throw new Error('独立 SAM 阶段使用 --input-dir，不接受模型请求或编辑参数');
    const result = await runSegmentation({...options, loaded});
    if (!result.dryRun) console.log(`SAM 分割${result.reused ? '复用完成' : '完成'}：${path.join(result.out, 'report.html')}`);
    return result;
  }
  if (options.inputDir) throw new Error('--input-dir 仅用于 SAM 分割或部件补全阶段');
  const overrides = options.stage === 'understanding'
    ? {model: options.understandingModel, protocol: options.understandingProtocol, outputMode: options.outputMode, stream: options.stream}
    : {model: options.completionModel};
  const stage = resolveStage(loaded, options.stage, overrides, {requireKey: !options.dryRun && !options.partsFile});
  if (options.listModels) {
    if (options.dryRun) throw new Error('--list-models 会请求模型列表，不能与 --dry-run 同时使用');
    const {payload} = await requestJson(stage, 'models', () => ({method: 'GET'}));
    if (!Array.isArray(payload.data)) throw new Error('模型列表响应缺少 data 数组，请核对中转站协议');
    console.log(payload.data.map((item) => item.id).filter((id) => typeof id === 'string').join('\n'));
    return;
  }
  if (options.dryRun) {
    const understanding = resolveStage(loaded, 'understanding', {model: options.understandingModel, protocol: options.understandingProtocol, outputMode: options.outputMode, stream: options.stream}, {requireKey: false});
    const completion = resolveStage(loaded, 'completion', {model: options.completionModel}, {requireKey: false});
    if (options.reference) readImage(options.reference, {pngOnly: true});
    for (const file of options.imagePaths) readImage(file, {pngOnly: true});
    console.log(JSON.stringify({understanding: publicStage(understanding), completion: publicStage(completion), sam: samInfo(loaded), credentials_present: Boolean(stage.apiKey && stage.baseUrl), network_requested: false}, null, 2));
    return;
  }
  if (!options.outDir) throw new Error('需要 --out-dir，建议每个实验使用独立目录');
  let result;
  if (options.stage === 'understanding') {
    if (!options.reference) throw new Error('结构理解需要 --reference');
    if (options.imagePaths.length || options.maskPath || options.promptFile) throw new Error('结构理解不接受图片编辑参数');
    console.log(options.partsFile ? '角色结构理解：离线计划几何校验' : `角色结构理解：${stage.model} (${stage.protocol})`);
    result = await runUnderstanding({...options, loaded, stage});
    console.log(`${result.reused ? '复用已校验产物' : '结构理解完成，语义待复核'}：${path.join(result.out, 'report.html')}`);
    if (options.stopAfter === 'segment') {
      result = await runSegmentation({...options, loaded, inputDir: result.out, outDir: path.join(result.out, 'segmentation')});
      console.log(`SAM 分割完成：${path.join(result.out, 'report.html')}`);
    }
  } else {
    if (options.reference) throw new Error('图片编辑使用 --image 指定局部画布；额外参考图也通过重复 --image 指定');
    console.log(`图片编辑接入测试：${stage.model}`);
    result = await runCompletionProbe({...options, stage});
    console.log(`${result.reused ? '复用已校验产物' : '原始候选已保存，尚未生成最终部件'}：${result.out}`);
  }
  return result;
}

if (process.argv[1] && path.resolve(process.argv[1]) === fileURLToPath(import.meta.url)) {
  main().catch((error) => { console.error(`模型阶段失败：${error.message}`); process.exitCode = 1; });
}