server.js

/**
 * New NAI - Node.js/Express 服务
 * 目标：保持与现有 FastAPI 接口一致，使前端无需改动即可运行。
 *
 * 实现接口：
 * - GET  /api/health
 * - GET  /api/config
 * - PUT  /api/config
 * - GET  /api/select-output-dir   （本地桌面环境：弹出目录选择器；无 GUI 环境会失败）
 * - POST /api/open-dir            （打开目录，若未传 path 则读取配置中的 output_dir）
 *
 * 静态资源：
 * - /       -> ./frontend （index.html 单页）
 * - /ring   -> ./ring     （提示音目录）
 *   - 兼容别名：/ring/ring.mp3 若实际不存在则回退到 new-notification-3-398649.mp3
 *
 * 注意：本文件尚未实现 /api/generate/* 生成接口；将在后续步骤补充。
 */

const path = require('path');
const fs = require('fs');
const fse = require('fs-extra');
const os = require('os');
const http = require('http');
const express = require('express');
const cors = require('cors');
const helmet = require('helmet');
const compression = require('compression');
const morgan = require('morgan');
const { spawn } = require('child_process');
const { generateT2I, generateI2I, generateInpaint } = require('./novelai');

const app = express();

// ---------- 常量与路径 ----------
const ROOT = __dirname;
const FRONTEND_DIR = path.join(ROOT, 'frontend');
const RING_DIR = path.join(ROOT, 'ring');
const IMAGE_DIR = path.join(ROOT, 'image');
const IMAGE2_DIR = path.join(ROOT, 'image2');
const CONFIG_PATH = path.join(ROOT, 'backend', 'config.json');

// ---------- 中间件 ----------
app.disable('x-powered-by');
app.use(cors({
  origin: '*',
  credentials: false,
  methods: ['GET', 'POST', 'PUT', 'DELETE', 'OPTIONS'],
  allowedHeaders: ['Content-Type', 'Authorization', 'x-api-key']
}));

// Helmet 安全头（简化版 CSP，允许被同源 iframe 嵌入；可按需扩展）
app.use(helmet({
  contentSecurityPolicy: {
    useDefaults: true,
    directives: {
      "default-src": ["'self'"],
      "img-src": ["'self'", "data:", "blob:"],
      "style-src": ["'self'", "'unsafe-inline'"],
      "script-src": ["'self'"],
      "font-src": ["'self'", "data:"],
      // 允许在 HF 页面中嵌入，以及前端与同源/HF 域通信
      "connect-src": ["'self'", "https://*.hf.space", "https://huggingface.co"],
      "frame-ancestors": ["'self'", "https://*.hf.space", "https://huggingface.co"],
    }
  },
  crossOriginResourcePolicy: { policy: "same-origin" }
}));

app.use(compression());
app.use(express.json({ limit: '50mb' })); // 前端传参与图片 Base64，放宽到 50MB
app.use(morgan('dev'));

// ---------- 工具函数：配置读写与默认值 ----------
const DEFAULT_CONFIG = {
  key: null,
  model: "nai-diffusion-3",
  sampler: "k_euler",
  steps: 28,
  scale: 5.0,
  cfg_rescale: 0.0,
  noise_schedule: "karras",
  uc_preset: 4,
  quality_toggle: true,
  legacy_uc: false,
  port: 7860,
  save_output: true,
  output_dir: path.join(ROOT, 'output'),
  // 提示音配置
  sound_enabled: false,
  sound_url: "/ring/ring.mp3"
};
let CURRENT_CFG = null; // HF: 内存态配置，避免写盘触发重启

function readConfig() {
  // HF 环境：优先使用内存态，避免文件写入引发 nodemon/watchdog 重启
  if (CURRENT_CFG) return { ...CURRENT_CFG };
  try {
    if (fs.existsSync(CONFIG_PATH)) {
      const raw = fs.readFileSync(CONFIG_PATH, 'utf-8');
      const fileCfg = JSON.parse(raw || '{}');
      CURRENT_CFG = { ...DEFAULT_CONFIG, ...fileCfg };
      return { ...CURRENT_CFG };
    }
  } catch (e) {
    // ignore and fallback
  }
  CURRENT_CFG = { ...DEFAULT_CONFIG };
  return { ...CURRENT_CFG };
}

function writeConfig(cfg) {
  // HF 环境：仅更新内存态，不写盘，避免引发容器内重启/热更新
  const merged = { ...DEFAULT_CONFIG, ...(cfg || {}) };
  CURRENT_CFG = merged;
  return { ...merged };
}

// ---------- 工具函数：系统命令 ----------
function runSpawn(cmd, args, options = {}) {
  return new Promise((resolve, reject) => {
    const child = spawn(cmd, args, {
      stdio: ['ignore', 'pipe', 'pipe'],
      shell: false,
      ...options
    });
    let stdout = '';
    let stderr = '';
    child.stdout.on('data', (d) => { stdout += d.toString(); });
    child.stderr.on('data', (d) => { stderr += d.toString(); });
    child.on('error', reject);
    child.on('close', (code) => {
      if (code === 0) resolve({ stdout, stderr, code });
      else reject(new Error(stderr || `Exit code ${code}`));
    });
  });
}

function sleep(ms) {
  return new Promise(r => setTimeout(r, ms));
}

// Windows: 使用 COM Shell.Application 弹出目录选择器
async function winBrowseForFolder() {
  // 方案1：PowerShell COM
  try {
    const psCmd = `$f=(New-Object -ComObject Shell.Application).BrowseForFolder(0,"选择保存目录",0); if($f){$f.Self.Path}`;
    const { stdout } = await runSpawn('powershell.exe', ['-NoProfile', '-Command', psCmd], { timeout: 60000 });
    const out = (stdout || '').trim();
    if (out && out.length > 0) return out;
  } catch (e) {
    // 继续尝试下一个方案
  }

  // 方案2：VBScript 兜底
  try {
    const vbs = [
      'Set sh = CreateObject("Shell.Application")',
      'Set f = sh.BrowseForFolder(0, "选择保存目录", 0)',
      'If (Not f Is Nothing) Then',
      '  WScript.Echo f.Self.Path',
      'End If'
    ].join('\n');

    const osTmp = os.tmpdir();
    const vbsPath = path.join(osTmp, `browse_${Date.now()}.vbs`);
    fs.writeFileSync(vbsPath, vbs, 'utf-8');
    try {
      const { stdout } = await runSpawn('cscript.exe', ['//nologo', vbsPath], { timeout: 60000 });
      const out = (stdout || '').trim();
      if (out && out.length > 0) return out;
    } finally {
      try { fs.unlinkSync(vbsPath); } catch {}
    }
  } catch (e) {
    // 继续尝试下一个方案
  }

  throw new Error('Windows 目录选择失败：所有方案均失败');
}

async function darwinChooseFolder() {
  const script = 'tell application "System Events" to POSIX path of (choose folder with prompt "选择保存目录")';
  const { stdout } = await runSpawn('osascript', ['-e', script], { timeout: 60000 });
  const out = (stdout || '').trim();
  if (out) return out;
  throw new Error('macOS 目录选择失败');
}

async function linuxZenity() {
  const { stdout } = await runSpawn('zenity', ['--file-selection', '--directory', '--title=选择保存目录'], { timeout: 60000 });
  const out = (stdout || '').trim();
  if (out) return out;
  throw new Error('Linux 目录选择失败 / zenity 不可用');
}

async function pickDirectory() {
  const sys = os.platform(); // win32, darwin, linux
  if (sys === 'win32') {
    return await winBrowseForFolder();
  } else if (sys === 'darwin') {
    return await darwinChooseFolder();
  } else if (sys === 'linux') {
    return await linuxZenity();
  }
  throw new Error(`不支持的系统平台：${sys}`);
}

async function openDirectory(p) {
  const sys = os.platform();
  const absPath = path.resolve(p);
  
  // 确保目录存在
  await fse.ensureDir(absPath);
  
  if (sys === 'win32') {
    // Windows: 使用 explorer
    try {
      await runSpawn('explorer.exe', [absPath]);
    } catch (e) {
      // 回退方案：使用 cmd start
      spawn('cmd', ['/c', 'start', '', absPath], { detached: true, stdio: 'ignore' }).unref();
    }
  } else if (sys === 'darwin') {
    await runSpawn('open', [absPath]);
  } else {
    // Linux
    await runSpawn('xdg-open', [absPath]);
  }
}

// ---------- API 路由 ----------

// 健康检查
app.get('/api/health', (req, res) => {
  res.json({ status: 'ok' });
});

// 获取配置
app.get('/api/config', (req, res) => {
  try {
    const cfg = readConfig();
    res.json(cfg);
  } catch (e) {
    res.status(500).json({ detail: String(e && e.message || e) });
  }
});

// 更新配置（部分字段）
app.put('/api/config', (req, res) => {
  try {
    const current = readConfig();
    const body = req.body || {};
    // 仅合并非 undefined 的字段；允许 null 写入（表示清空）
    const next = { ...current };
    for (const k of Object.keys(body)) {
      next[k] = body[k];
    }
    const saved = writeConfig(next);
    res.json(saved);
  } catch (e) {
    res.status(500).json({ detail: String(e && e.message || e) });
  }
});

// 选择输出目录（桌面环境）
app.get('/api/select-output-dir', async (req, res) => {
  // HF 环境不支持系统目录选择，直接返回 501，前端已做兜底提示
  return res.status(501).json({ detail: 'HF 环境不支持目录选择，请在“保存目录”中手动填写或留空。' });
});

// 打开目录
app.post('/api/open-dir', async (req, res) => {
  // HF 环境不支持打开系统目录，统一返回 501
  return res.status(501).json({ detail: 'HF 环境不支持打开系统目录（不影响生成与下载）。' });
});

// 生成接口：T2I
app.post('/api/generate/t2i', async (req, res) => {
 try {
   const cfg = readConfig();
   if (!cfg.key) {
     return res.status(400).json({ detail: '尚未配置 key，请先在配置中设置 key。' });
   }
   const b = req.body || {};
   const { dataUri, savedPath } = await generateT2I(cfg, {
     prompt: b.prompt,
     negative: b.negative || '',
     width: b.width ?? 768,
     height: b.height ?? 768,
     scale: b.scale ?? null,
     steps: b.steps ?? null,
     sampler: b.sampler ?? null,
     noise_schedule: b.noise_schedule ?? null,
     seed: b.seed ?? -1,
     variety: !!b.variety,
     decrisp: !!b.decrisp,
     cfg_rescale: b.cfg_rescale ?? null,
   });
   return res.json({ image_base64: dataUri, saved_path: savedPath });
 } catch (e) {
   return res.status(500).json({ detail: String((e && e.message) || e) });
 }
});

// 生成接口：I2I
app.post('/api/generate/i2i', async (req, res) => {
 try {
   const cfg = readConfig();
   if (!cfg.key) {
     return res.status(400).json({ detail: '尚未配置 key，请先在配置中设置 key。' });
   }
   const b = req.body || {};
   const { dataUri, savedPath } = await generateI2I(cfg, {
     positive: b.positive || '',
     negative: b.negative || '',
     image_base64: b.image_base64,
     width: b.width ?? null,
     height: b.height ?? null,
     scale: b.scale ?? null,
     steps: b.steps ?? null,
     sampler: b.sampler ?? null,
     noise_schedule: b.noise_schedule ?? null,
     strength: b.strength ?? 0.5,
     noise: b.noise ?? 0.0,
     seed: b.seed ?? -1,
     variety: !!b.variety,
     decrisp: !!b.decrisp,
     cfg_rescale: b.cfg_rescale ?? null,
   });
   return res.json({ image_base64: dataUri, saved_path: savedPath });
 } catch (e) {
   return res.status(500).json({ detail: String((e && e.message) || e) });
 }
});

// 生成接口：Inpaint
app.post('/api/generate/inpaint', async (req, res) => {
 try {
   const cfg = readConfig();
   if (!cfg.key) {
     return res.status(400).json({ detail: '尚未配置 key，请先在配置中设置 key。' });
   }
   const b = req.body || {};
   const { dataUri, savedPath } = await generateInpaint(cfg, {
     positive: b.positive || '',
     negative: b.negative || '',
     image_base64: b.image_base64,
     mask_base64: b.mask_base64,
     add_original_image: !!b.add_original_image,
     width: b.width ?? null,
     height: b.height ?? null,
     scale: b.scale ?? null,
     steps: b.steps ?? null,
     sampler: b.sampler ?? null,
     noise_schedule: b.noise_schedule ?? null,
     strength: b.strength ?? 0.5,
     noise: b.noise ?? 0.0,
     seed: b.seed ?? -1,
     variety: !!b.variety,
     decrisp: !!b.decrisp,
     cfg_rescale: b.cfg_rescale ?? null,
   });
   return res.json({ image_base64: dataUri, saved_path: savedPath });
 } catch (e) {
   return res.status(500).json({ detail: String((e && e.message) || e) });
 }
});

// favicon 占位，避免某些浏览器 404 导致干扰
app.get('/favicon.ico', (req, res) => res.status(204).end());

// ---------- 静态资源 ----------

// 背景图片目录（电脑端）
app.use('/image', express.static(IMAGE_DIR, { fallthrough: true }));

// 背景图片目录（手机端）
app.use('/image2', express.static(IMAGE2_DIR, { fallthrough: true }));

// 铃声目录
app.use('/ring', express.static(RING_DIR, { fallthrough: true }));

// 兼容别名：/ring/ring.mp3 -> 若 ring.mp3 不存在则回退到 new-notification-3-398649.mp3
app.get('/ring/ring.mp3', (req, res, next) => {
  const main = path.join(RING_DIR, 'ring.mp3');
  const alt = path.join(RING_DIR, 'new-notification-3-398649.mp3');
  if (fs.existsSync(main)) return res.sendFile(main);
  if (fs.existsSync(alt)) return res.sendFile(alt);
  return res.status(404).end();
});

// 前端静态资源（单页）
app.use('/', express.static(FRONTEND_DIR, { index: 'index.html' }));

// ---------- 启动服务 ----------
function resolvePort() {
  // HF Docker 要求监听 $PORT（若无则回退 7860）
  const envPort = parseInt(process.env.PORT || '', 10);
  if (!Number.isNaN(envPort) && envPort > 0) return envPort;
  return 7860;
}

const PORT = resolvePort();
// HF 必须使用 0.0.0.0 以便外部访问
const HOST = process.env.HOST || '0.0.0.0';

const server = http.createServer(app);
server.listen(PORT, HOST, () => {
  // HF 环境显示 0.0.0.0，但实际通过 Space URL 访问
  const displayUrl = HOST === '0.0.0.0' ? `http://0.0.0.0:${PORT}` : `http://${HOST}:${PORT}`;
  
  // HF 环境不自动打开浏览器
  const shouldOpen = process.env.AUTO_OPEN_BROWSER !== '0' && HOST === '127.0.0.1';
  if (shouldOpen) {
    setTimeout(() => {
      try {
        const localUrl = `http://127.0.0.1:${PORT}`;
        const platform = os.platform();
        if (platform === 'win32') {
          spawn('cmd', ['/c', 'start', '', localUrl], { detached: true, stdio: 'ignore' }).unref();
        } else if (platform === 'darwin') {
          spawn('open', [localUrl], { detached: true, stdio: 'ignore' }).unref();
        } else {
          spawn('xdg-open', [localUrl], { detached: true, stdio: 'ignore' }).unref();
        }
      } catch { /* ignore */ }
    }, 1500);
  }
  
  // eslint-disable-next-line no-console
  console.log(`[New NAI HF] 服务运行于 ${displayUrl}`);
  console.log(`[New NAI HF] ${HOST === '0.0.0.0' ? 'HF Space 通过公共 URL 访问' : '按 Ctrl+C 停止服务'}`);
});

// 导出 app 以便测试或其他入口使用
module.exports = { app };