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();

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

    const baseUrl = aiConfig.ai_base_url || 'https://api.deepseek.com';
    const response = await fetch(`${baseUrl}/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 || '';
  }

  // Helper: generate title for a note
  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;
  }

  // Helper: generate content for a note
  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(t => t.replace('#', '').trim());
    return { content: generated.trim(), topics };
  }

  // Helper: generate title and content together (legacy fallback)
  async function generateText(productTitle) {
    const title = await generateTitle(productTitle);
    const { content, topics } = await generateContent(productTitle);
    return { title, content, topics };
  }

  async function createNoteForProduct(product, options = {}) {
    // Dedup: if a note is already being generated for this product, skip
    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) {
      console.log(`[generate] Note already generating for product ${product.id}, reusing ${existing.id}`);
      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 };
  }

  // Generate note for a product using AI
  router.post('/note', async (req, res) => {
    const { product_id, shop_id } = 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' });

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

    try {
      const result = await createNoteForProduct(product);
      res.json({ note_id: result.noteId, title: result.title, content: result.content, topics: result.topics, image_paths: result.imagePaths });
    } catch (error) {
      db.prepare("UPDATE notes SET status='failed', error_message=?, updated_at=CURRENT_TIMESTAMP WHERE id=(SELECT id FROM notes WHERE product_id=? ORDER BY created_at DESC LIMIT 1)")
        .run(error.message, product_id);
      db.prepare("UPDATE products SET status='pending', updated_at=CURRENT_TIMESTAMP WHERE id=?")
        .run(product_id);
      res.status(500).json({ error: error.message });
    }
  });

  // Generate image for a note
  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, 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 results = [];
    for (const product of products) {
      try {
        db.prepare("UPDATE products SET status = 'generating', updated_at = CURRENT_TIMESTAMP WHERE id = ?").run(product.id);
        const result = await createNoteForProduct(product, { image_count, image_size, image_prompt_template });
        results.push({ product_id: product.id, note_id: result.noteId, title: result.title, images: result.imagePaths.length });
      } catch (error) {
        results.push({ product_id: product.id, error: error.message });
        db.prepare("UPDATE products SET status = 'pending', updated_at = CURRENT_TIMESTAMP WHERE id = ?").run(product.id);
      }
    }

    return res.json({ total: products.length, results, images_generated: results.reduce((sum, item) => sum + (item.images || 0), 0) });
  }

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

  return router;
}

module.exports = createGenerateRoutes;