xiaohongshufabu/server/routes/generate.js

const express = require('express');
const { v4: uuidv4 } = require('uuid');
const { generateImage } = require('../lib/generate-image');

function createGenerateRoutes(db) {
  const router = express.Router();
  const runningJobs = new Set();

  async function callAI(prompt, maxTokens = 2000) {
    const aiConfig = {};
    const keys = ['ai_api_key', 'ai_base_url', 'ai_model', 'ai_temperature'];
    for (const key of keys) {
      const row = db.prepare('SELECT value FROM configs WHERE key = ?').get(key);
      aiConfig[key] = row?.value || '';
    }
    if (!aiConfig.ai_api_key) throw new Error('AI API key not configured. Please set it in Settings.');

    const response = await fetch(`${aiConfig.ai_base_url || 'https://api.deepseek.com'}/chat/completions`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${aiConfig.ai_api_key}`,
      },
      body: JSON.stringify({
        model: aiConfig.ai_model || 'deepseek-chat',
        messages: [{ role: 'user', content: prompt }],
        temperature: parseFloat(aiConfig.ai_temperature || '0.8'),
        max_tokens: maxTokens,
      }),
    });

    if (!response.ok) {
      const err = await response.text();
      throw new Error(`AI API error: ${response.status} ${err}`);
    }

    const data = await response.json();
    return data.choices?.[0]?.message?.content || '';
  }

  async function generateTitle(productTitle) {
    const titleRow = db.prepare('SELECT value FROM configs WHERE key = ?').get('title_prompt_template');
    const fallbackRow = db.prepare('SELECT value FROM configs WHERE key = ?').get('prompt_template');
    const template = titleRow?.value || fallbackRow?.value || '为商品写一个小红书标题：{{title}}';
    const prompt = template.replace(/\{\{title\}\}/g, productTitle);
    const generated = await callAI(prompt, 500);
    const title = generated.replace(/^#+\s*/, '').replace(/[\n\r]/g, '').trim();
    return title || productTitle;
  }

  async function generateContent(productTitle) {
    const contentRow = db.prepare('SELECT value FROM configs WHERE key = ?').get('content_prompt_template');
    const fallbackRow = db.prepare('SELECT value FROM configs WHERE key = ?').get('prompt_template');
    const template = contentRow?.value || fallbackRow?.value || '为商品写一篇小红书正文：{{title}}';
    const prompt = template.replace(/\{\{title\}\}/g, productTitle);
    const generated = await callAI(prompt, 2000);
    const topicMatches = generated.match(/#[^\s#]+/g) || [];
    const topics = topicMatches.map((topic) => topic.replace('#', '').trim());
    return { content: generated.trim(), topics };
  }

  async function generateText(productTitle) {
    const title = await generateTitle(productTitle);
    const { content, topics } = await generateContent(productTitle);
    return { title, content, topics };
  }

  async function createNoteForProduct(product, options = {}) {
    const existing = db.prepare(
      "SELECT id, title, content, topics, image_paths, image_prompt FROM notes WHERE product_id = ? AND status = 'generating' ORDER BY created_at DESC LIMIT 1"
    ).get(product.id);
    if (existing) {
      return {
        noteId: existing.id,
        title: existing.title,
        content: existing.content,
        topics: JSON.parse(existing.topics || '[]'),
        imagePaths: JSON.parse(existing.image_paths || '[]'),
        imagePrompt: existing.image_prompt || '',
      };
    }

    const noteId = uuidv4();
    db.prepare("INSERT INTO notes (id, product_id, shop_id, title, content, status) VALUES (?, ?, ?, ?, '', 'generating')")
      .run(noteId, product.id, product.shop_id, product.title);

    const { title, content, topics } = await generateText(product.title);
    const imageResult = await generateImage(db, noteId, product.title, options);
    const imagePaths = imageResult.imagePaths || [];
    const imagePrompt = imageResult.imagePrompt || '';

    db.prepare(
      "UPDATE notes SET title=?, content=?, topics=?, image_paths=?, image_prompt=?, status='generated', updated_at=CURRENT_TIMESTAMP WHERE id=?"
    ).run(title, content, JSON.stringify(topics), JSON.stringify(imagePaths), imagePrompt, noteId);
    db.prepare("UPDATE products SET status='generated', updated_at=CURRENT_TIMESTAMP WHERE id=?").run(product.id);

    return { noteId, title, content, topics, imagePaths, imagePrompt };
  }

  function updateJob(jobId, fields) {
    const keys = Object.keys(fields);
    if (keys.length === 0) return;
    const setSql = keys.map((key) => `${key} = ?`).join(', ');
    const values = keys.map((key) => fields[key]);
    values.push(jobId);
    db.prepare(`UPDATE generation_jobs SET ${setSql}, updated_at=CURRENT_TIMESTAMP WHERE id = ?`).run(...values);
  }

  async function runGenerationJob(jobId, productIds, countPerProduct, options = {}) {
    if (runningJobs.has(jobId)) return;
    runningJobs.add(jobId);
    try {
      updateJob(jobId, { status: 'running', message: '开始生成' });
      let completedNotes = 0;
      let failedNotes = 0;

      for (const productId of productIds) {
        const product = db.prepare('SELECT * FROM products WHERE id = ?').get(productId);
        if (!product) {
          failedNotes += countPerProduct;
          completedNotes += countPerProduct;
          updateJob(jobId, {
            completed_notes: completedNotes,
            failed_notes: failedNotes,
            current_product_id: productId,
            current_product_title: '',
            message: '商品不存在',
          });
          continue;
        }

        db.prepare("UPDATE products SET status='generating', updated_at=CURRENT_TIMESTAMP WHERE id=?").run(product.id);
        updateJob(jobId, {
          current_product_id: product.id,
          current_product_title: product.title,
          message: `正在生成：${product.title}`,
        });

        for (let index = 0; index < countPerProduct; index += 1) {
          try {
            await createNoteForProduct(product, options);
          } catch (error) {
            failedNotes += 1;
          } finally {
            completedNotes += 1;
            updateJob(jobId, {
              completed_notes: completedNotes,
              failed_notes: failedNotes,
              message: `进度 ${completedNotes}/${productIds.length * countPerProduct}`,
            });
          }
        }

        db.prepare("UPDATE products SET status='generated', updated_at=CURRENT_TIMESTAMP WHERE id=?").run(product.id);
      }

      updateJob(jobId, {
        status: 'completed',
        message: failedNotes > 0 ? `完成，失败 ${failedNotes} 条` : '全部完成',
      });
    } catch (error) {
      updateJob(jobId, { status: 'failed', message: error.message });
    } finally {
      runningJobs.delete(jobId);
    }
  }

  function enqueueGeneration(jobId, productIds, countPerProduct, options) {
    setImmediate(() => {
      runGenerationJob(jobId, productIds, countPerProduct, options);
    });
  }

  router.get('/jobs', (req, res) => {
    const { shop_id } = req.query;
    let sql = 'SELECT * FROM generation_jobs WHERE 1=1';
    const params = [];
    if (shop_id) {
      sql += ' AND shop_id = ?';
      params.push(shop_id);
    }
    sql += ' ORDER BY created_at DESC LIMIT 20';
    res.json(db.prepare(sql).all(...params));
  });

  router.post('/note', (req, res) => {
    const { product_id, shop_id, count_per_product, image_count, image_size, image_prompt_template } = req.body;
    if (!product_id || !shop_id) {
      return res.status(400).json({ error: 'product_id and shop_id required' });
    }

    const product = db.prepare('SELECT * FROM products WHERE id = ?').get(product_id);
    if (!product) return res.status(404).json({ error: 'Product not found' });

    const repeatCount = Math.max(1, parseInt(count_per_product || '1', 10) || 1);
    const jobId = uuidv4();
    db.prepare(
      'INSERT INTO generation_jobs (id, shop_id, status, total_products, total_notes, message) VALUES (?, ?, ?, ?, ?, ?)'
    ).run(jobId, shop_id, 'pending', 1, repeatCount, '等待开始');
    enqueueGeneration(jobId, [product_id], repeatCount, { image_count, image_size, image_prompt_template });
    res.status(201).json({ job_id: jobId, status: 'pending' });
  });

  router.post('/image', async (req, res) => {
    const { note_id } = req.body;
    if (!note_id) return res.status(400).json({ error: 'note_id required' });

    const note = db.prepare('SELECT * FROM notes WHERE id = ?').get(note_id);
    if (!note) return res.status(404).json({ error: 'Note not found' });

    try {
      const imageResult = await generateImage(db, note_id, note.title);
      const imagePaths = imageResult.imagePaths || [];
      const existingPaths = JSON.parse(note.image_paths || '[]');
      const allPaths = [...existingPaths, ...imagePaths];
      db.prepare("UPDATE notes SET image_paths=?, image_prompt=?, updated_at=CURRENT_TIMESTAMP WHERE id=?")
        .run(JSON.stringify(allPaths), imageResult.imagePrompt || note.image_prompt || '', note_id);
      res.json({ image_paths: imagePaths, total: allPaths.length });
    } catch (error) {
      res.status(500).json({ error: error.message });
    }
  });

  async function batchGenerate(req, res) {
    const { shop_id, product_ids, count_per_product, image_count, image_size, image_prompt_template } = req.body;
    if (!shop_id) return res.status(400).json({ error: 'shop_id required' });

    let sql = 'SELECT * FROM products WHERE shop_id = ?';
    const params = [shop_id];
    if (Array.isArray(product_ids) && product_ids.length > 0) {
      sql += ` AND id IN (${product_ids.map(() => '?').join(',')})`;
      params.push(...product_ids);
    }
    const products = db.prepare(sql).all(...params);
    if (products.length === 0) return res.json({ count: 0, message: 'No pending products' });

    const repeatCount = Math.max(1, parseInt(count_per_product || '1', 10) || 1);
    const jobId = uuidv4();
    db.prepare(
      'INSERT INTO generation_jobs (id, shop_id, status, total_products, total_notes, message) VALUES (?, ?, ?, ?, ?, ?)'
    ).run(jobId, shop_id, 'pending', products.length, products.length * repeatCount, '等待开始');
    enqueueGeneration(jobId, products.map((product) => product.id), repeatCount, { image_count, image_size, image_prompt_template });
    res.status(201).json({ job_id: jobId, total: products.length, total_notes: products.length * repeatCount });
  }

  router.post('/', batchGenerate);
  router.post('/batch', batchGenerate);

  return router;
}

module.exports = createGenerateRoutes;