hyreal-gen-pic/server/routes/generate.js

'use strict';

// 视觉生成路由：品牌资产/菜品图生成提交、任务查询、任务中心、提示词配置、画板保存、打开目录

const fs = require('fs');
const path = require('path');
const { execFile } = require('child_process');
const { sendJson, readBody, sanitizeFileName } = require('../util.js');
const { getImageSize } = require('../../imageMeta.js');

const POSTER_CANVAS_SPECS = {
  banner: { fileName: '店内海报_1138x292.png', width: 1138, height: 292 },
  signage: { fileName: '招牌头图_750x288.png', width: 750, height: 288 },
};

async function handleGenerateBrand(ctx, body, res) {
  const { root, folderName, shopName, kind, prompt, size, aspectRatio, refDishFiles, addLogo, logoPosition, logoMode, model } = body;
  const outNameMap = { logo: 'logo', banner: 'banner', signage: '招牌头图', sticker: '贴纸' };
  if (!root || !folderName || !kind || !outNameMap[kind] || !prompt) {
    return sendJson(res, 400, { error: '参数不完整' });
  }
  const shopDir = ctx.workspace.safeShopDir(root, folderName);
  if (!shopDir) return sendJson(res, 403, { error: '店铺目录无效' });
  // logo/sticker are themselves brand-mark assets, adding the shop logo onto them is not applicable
  const shouldAddLogo = !!addLogo && kind !== 'logo';
  const { images, prompt: finalPrompt, warning, compositeOverlays } = ctx.generation.buildImagesAndPrompt(shopDir, prompt, refDishFiles, {
    addLogo: shouldAddLogo,
    logoPosition,
    logoMode,
  });
  const job = ctx.generation.createJob({
    kind,
    root,
    folderName,
    shopName,
    model,
    prompt: finalPrompt,
    size: size || '1K',
    aspectRatio: aspectRatio || undefined,
    images,
    outputDir: path.join(shopDir, '_生成图片'),
    outputBaseName: outNameMap[kind],
    warning,
    compositeOverlays,
  });
  sendJson(res, 200, { jobId: job.id });
}

async function handleGenerateDish(ctx, body, res) {
  const { root, folderName, shopName, dishFile, dishName, prompt, size, addLogo, logoPosition, logoMode, addSticker, stickerPosition, stickerMode, model } = body;
  if (!root || !folderName || !dishFile || !prompt) {
    return sendJson(res, 400, { error: '参数不完整' });
  }
  const shopDir = ctx.workspace.safeShopDir(root, folderName);
  if (!shopDir) return sendJson(res, 403, { error: '店铺目录无效' });
  const dishPath = path.resolve(shopDir, dishFile);
  if (!ctx.workspace.isWorkspacePath(dishPath) || !dishPath.startsWith(shopDir + path.sep)) {
    return sendJson(res, 403, { error: '菜品文件无效' });
  }
  let dishImage;
  try {
    dishImage = ctx.generation.readAsDataUri(dishPath);
  } catch (e) {
    return sendJson(res, 400, { error: '读取菜品原图失败: ' + e.message });
  }
  const { images: overlayImages, prompt: finalPrompt, warning, compositeOverlays } = ctx.generation.buildImagesAndPrompt(shopDir, prompt, [], {
    addLogo: !!addLogo,
    logoPosition,
    logoMode,
    addSticker: !!addSticker,
    stickerPosition,
    stickerMode,
    sourceImageFirst: true,
  });
  const job = ctx.generation.createJob({
    kind: 'dish',
    root,
    folderName,
    shopName,
    dishName,
    model,
    prompt: finalPrompt,
    size: size || '1K',
    aspectRatio: '1:1',
    images: [dishImage, ...overlayImages],
    outputDir: path.join(shopDir, '_生成图片', '菜品图'),
    outputBaseName: sanitizeFileName(dishName),
    warning,
    compositeOverlays,
  });
  sendJson(res, 200, { jobId: job.id });
}

function publicTask(job) {
  return {
    id: job.id,
    type: 'generation',
    status: job.state,
    root: job.root,
    folderName: job.folderName,
    shopName: job.shopName,
    dishName: job.dishName,
    name: job.dishName || job.kind,
    kind: job.kind,
    model: job.model,
    providerTaskId: job.taskId,
    progress: job.progress,
    cost: job.cost,
    resultUrl: job.resultUrl,
    remoteUrl: job.remoteUrl,
    width: job.width,
    height: job.height,
    fileSizeBytes: job.fileSizeBytes,
    error: job.error,
    warning: job.warning,
    createdAt: job.createdAt,
    startedAt: job.startedAt,
    finishedAt: job.finishedAt,
    updatedAt: job.updatedAt || job.createdAt,
  };
}

// 采集任务映射为任务中心条目（与生成任务并列展示）
function crawlTaskFromState(cs) {
  if (!cs || cs.status === 'idle') return null;
  const shops = Array.isArray(cs.shops) ? cs.shops : [];
  const done = shops.filter((s) => s.state === 'done' || s.state === 'skipped_completed').length;
  const failedShops = shops.filter((s) => s.state === 'failed').length;
  const statusMap = { running: 'running', done: 'success', stopped: 'failed', interrupted: 'failed', error: 'failed' };
  const labelMap = { running: '采集中', done: '采集完成', stopped: '采集已停止', interrupted: '采集被中断', error: '采集出错' };
  const pct = shops.length ? Math.round((done / shops.length) * 100) : 0;
  return {
    id: 'crawl-current',
    type: 'crawl',
    status: statusMap[cs.status] || 'running',
    name: `门店采集 · ${labelMap[cs.status] || cs.status}`,
    kind: 'crawl',
    shopName: '',
    model: '',
    progress: pct + '%',
    detail: `${done}/${shops.length} 家门店完成${failedShops ? ` · ${failedShops} 家失败` : ''}`,
    warning: cs.status === 'interrupted' ? '服务重启导致中断，未完成的门店可重新勾选采集' : '',
    createdAt: cs.startedAt,
    startedAt: cs.startedAt,
    finishedAt: cs.finishedAt,
    updatedAt: cs.finishedAt || cs.startedAt || Date.now(),
  };
}

function register(app, ctx) {
  app.post('/api/generate/brand', async ({ req, res }) => {
    const body = await readBody(req);
    return handleGenerateBrand(ctx, body, res);
  });

  app.post('/api/generate/dish', async ({ req, res }) => {
    const body = await readBody(req);
    return handleGenerateDish(ctx, body, res);
  });

  app.get('/api/job/:id', async ({ res, params }) => {
    const job = ctx.generation.getJob(params.id);
    if (!job) return sendJson(res, 404, { error: 'job not found' });
    sendJson(res, 200, {
      id: job.id,
      kind: job.kind,
      dishName: job.dishName,
      state: job.state,
      progress: job.progress,
      cost: job.cost,
      resultUrl: job.resultUrl,
      width: job.width,
      height: job.height,
      fileSizeBytes: job.fileSizeBytes,
      error: job.error,
      warning: job.warning,
    });
  });

  app.get('/api/tasks', async ({ res, query }) => {
    if (query.root && !ctx.workspace.isWorkspaceRoot(query.root)) {
      return sendJson(res, 403, { error: '仅允许访问当前输出根目录' });
    }
    const root = ctx.workspace.getWorkspaceRoot();
    const tasks = ctx.generation.listJobs()
      .filter((job) => path.resolve(job.root) === root)
      .sort((a, b) => b.createdAt - a.createdAt)
      .map(publicTask);
    const crawlTask = crawlTaskFromState(ctx.crawl.getState());
    if (crawlTask) tasks.unshift(crawlTask);
    const summary = tasks.reduce((acc, task) => {
      acc[task.status] = (acc[task.status] || 0) + 1;
      return acc;
    }, { all: tasks.length });
    sendJson(res, 200, { root, summary, tasks });
  });

  app.post('/api/save-config', async ({ req, res }) => {
    const body = await readBody(req);
    const { root, folderName, shopName, prompts } = body;
    if (!root || !folderName) return sendJson(res, 400, { error: '参数不完整' });
    const shopDir = ctx.workspace.safeShopDir(root, folderName);
    if (!shopDir) return sendJson(res, 403, { error: '店铺目录无效' });
    const dir = path.join(shopDir, '_生成图片');
    fs.mkdirSync(dir, { recursive: true });
    const cfgPath = path.join(dir, 'prompt-config.json');
    let existing = {};
    if (fs.existsSync(cfgPath)) {
      try {
        existing = JSON.parse(fs.readFileSync(cfgPath, 'utf8'));
      } catch {
        // start fresh if corrupt
      }
    }
    const merged = {
      shopName: shopName !== undefined ? shopName : existing.shopName,
      prompts: { ...(existing.prompts || {}), ...(prompts || {}) },
    };
    fs.writeFileSync(cfgPath, JSON.stringify(merged, null, 2), 'utf8');
    sendJson(res, 200, { ok: true });
  });

  app.post('/api/save-poster-canvas', async ({ req, res }) => {
    const body = await readBody(req);
    const { root, folderName, kind, imageDataUrl } = body;
    const spec = POSTER_CANVAS_SPECS[kind];
    if (!root || !folderName || !spec || !imageDataUrl) {
      return sendJson(res, 400, { error: '参数不完整' });
    }
    if (typeof folderName !== 'string' || folderName === '.' || folderName === '..' || /[\\/]/.test(folderName)) {
      return sendJson(res, 400, { error: '店铺目录名无效' });
    }
    const dataUrlMatch = typeof imageDataUrl === 'string' && /^data:image\/png;base64,([A-Za-z0-9+/=]+)$/.exec(imageDataUrl);
    if (!dataUrlMatch) {
      return sendJson(res, 400, { error: '画板图必须为 PNG 格式' });
    }

    const base64 = dataUrlMatch[1];
    if (!base64 || base64.length > 10 * 1024 * 1024) {
      return sendJson(res, 400, { error: '画板图数据无效或过大' });
    }

    let imageBuffer;
    try {
      imageBuffer = Buffer.from(base64, 'base64');
    } catch {
      return sendJson(res, 400, { error: '画板图数据无法读取' });
    }
    const pngSignature = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
    if (!imageBuffer.length || !imageBuffer.subarray(0, 8).equals(pngSignature)) {
      return sendJson(res, 400, { error: '画板图不是有效的 PNG 文件' });
    }
    const dimensions = getImageSize(imageBuffer);
    if (!dimensions || dimensions.width !== spec.width || dimensions.height !== spec.height) {
      return sendJson(res, 400, { error: `画板图尺寸必须为 ${spec.width}×${spec.height}` });
    }

    const resolvedRoot = path.resolve(root);
    if (!ctx.workspace.isWorkspaceRoot(resolvedRoot)) return sendJson(res, 403, { error: '仅允许访问当前输出根目录' });
    const shopDir = path.resolve(resolvedRoot, folderName);
    if (shopDir !== path.join(resolvedRoot, folderName) || !shopDir.startsWith(resolvedRoot + path.sep)) {
      return sendJson(res, 400, { error: '店铺目录无效' });
    }
    const dir = path.join(shopDir, '_生成图片');
    const savePath = path.join(dir, spec.fileName);
    try {
      fs.mkdirSync(dir, { recursive: true });
      fs.writeFileSync(savePath, imageBuffer);
    } catch (e) {
      return sendJson(res, 500, { error: '保存画板图失败: ' + e.message });
    }
    sendJson(res, 200, { ok: true, fileName: spec.fileName, width: dimensions.width, height: dimensions.height, fileSizeBytes: imageBuffer.length });
  });

  app.post('/api/open-folder', async ({ req, res }) => {
    const body = await readBody(req);
    if (body.root && !ctx.workspace.isWorkspaceRoot(body.root)) return sendJson(res, 403, { error: '仅允许访问当前输出根目录' });
    const dir = body.dir || path.join(body.root || ctx.workspace.getWorkspaceRoot(), body.folderName || '', '_生成图片');
    const fallbackDir = body.dir ? null : path.join(body.root || ctx.workspace.getWorkspaceRoot(), body.folderName || '');
    const targetDir = fs.existsSync(dir) ? dir : (fallbackDir && fs.existsSync(fallbackDir) ? fallbackDir : null);
    if (!targetDir || !ctx.workspace.isWorkspacePath(targetDir)) return sendJson(res, 400, { error: '目录不存在或不在输出根目录内: ' + dir });
    const command = process.platform === 'win32' ? 'explorer' : process.platform === 'darwin' ? 'open' : 'xdg-open';
    execFile(command, [targetDir], () => {});
    sendJson(res, 200, { ok: true, opened: targetDir });
  });
}

module.exports = { register };