require('dotenv').config();
const express = require('express');
const multer = require('multer');
const { S3Client, PutObjectCommand, DeleteObjectCommand } = require('@aws-sdk/client-s3');
const axios = require('axios');
const fs = require('fs');
const path = require('path');
const cors = require('cors');

// 导入新模块
const { planImages } = require('./lib/brain');
const { injectConstraints, getAspectRatio } = require('./lib/constraint-injector');

const app = express();
const port = process.env.PORT || 3000;

app.use(cors());
app.use(express.json({ limit: '50mb' }));
app.use(express.static('public'));

// Configure R2 Client
const bucketName = process.env.R2_BUCKET_NAME || 'ai-flow';
console.log('Using R2 Bucket:', bucketName);

const r2Client = new S3Client({
  region: 'auto',
  endpoint: `https://${process.env.R2_ACCOUNT_ID}.r2.cloudflarestorage.com`,
  credentials: {
    accessKeyId: process.env.R2_ACCESS_KEY_ID,
    secretAccessKey: process.env.R2_SECRET_ACCESS_KEY,
  },
  forcePathStyle: true,
});

// Multer for memory storage (upload directly to R2)
const upload = multer({ storage: multer.memoryStorage() });

// ============================================
// 原有API端点
// ============================================

// Endpoint to get prompt template
app.get('/api/prompt-template', (req, res) => {
  try {
    const promptPath = path.join(__dirname, 'prompt.md');
    const content = fs.readFileSync(promptPath, 'utf-8');
    res.json({ content });
  } catch (error) {
    res.status(500).json({ error: 'Failed to read prompt template' });
  }
});

// Endpoint to upload to R2
app.post('/api/upload-r2', upload.array('files'), async (req, res) => {
  try {
    const files = req.files;
    if (!files || files.length === 0) {
      return res.status(400).json({ error: 'No files uploaded' });
    }

    const uploadedUrls = [];

    for (const file of files) {
      const fileName = `${Date.now()}-${file.originalname}`;
      const command = new PutObjectCommand({
        Bucket: bucketName,
        Key: fileName,
        Body: file.buffer,
        ContentType: file.mimetype,
      });

      await r2Client.send(command);
      
      const publicUrl = process.env.R2_PUBLIC_DOMAIN 
        ? `${process.env.R2_PUBLIC_DOMAIN}/${fileName}`
        : `https://${process.env.R2_ACCOUNT_ID}.r2.cloudflarestorage.com/${bucketName}/${fileName}`;
        
      uploadedUrls.push(publicUrl);
    }

    res.json({ urls: uploadedUrls });
  } catch (error) {
    console.error('R2 Upload Error:', error);
    res.status(500).json({ error: 'Upload failed: ' + error.message });
  }
});

// Endpoint to delete from R2
app.delete('/api/delete-r2', async (req, res) => {
  try {
    const { url } = req.body;
    if (!url) return res.status(400).json({ error: 'URL is required' });

    const urlObj = new URL(url);
    const key = decodeURIComponent(urlObj.pathname.substring(1));

    const bucketName = process.env.R2_BUCKET_NAME || 'ai-flow';
    
    await r2Client.send(new DeleteObjectCommand({
      Bucket: bucketName,
      Key: key,
    }));

    res.json({ success: true });
  } catch (error) {
    console.error('R2 Delete Error:', error);
    res.status(500).json({ error: 'Delete failed: ' + error.message });
  }
});

// Endpoint to generate prompts (LLM) - 保留原有功能
app.post('/api/generate-prompts', async (req, res) => {
  try {
    const { prompt } = req.body;
    
    const response = await axios.post('https://api2img.shubiaobiao.com/v1/chat/completions', {
      model: 'gemini-3-pro-preview',
      messages: [
        { role: 'user', content: prompt }
      ],
      stream: false
    }, {
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${process.env.API_KEY}`
      }
    });

    res.json({ data: response.data });
  } catch (error) {
    console.error('LLM API Error:', error.response?.data || error.message);
    res.status(500).json({ error: 'Prompt generation failed' });
  }
});

// Endpoint to generate image (Sync) - 保留原有功能
app.post('/api/generate-image', async (req, res) => {
  try {
    const { prompt, aspectRatio, refImages } = req.body;
    console.log('Generate Image Request:', { promptLength: prompt?.length, refImagesCount: refImages?.length });
    
    let finalPrompt = prompt;
    let ar = aspectRatio || '1:1';
    
    const arMatch = prompt.match(/--ar\s+(\d+:\d+)/i);
    if (arMatch) {
      ar = arMatch[1];
    }

    const payload = {
      model: 'gemini-3-pro-image-preview',
      prompt: finalPrompt,
      n: 1,
      size: "1K",
      aspect_ratio: ar
    };

    if (refImages && Array.isArray(refImages) && refImages.length > 0) {
      payload.image_urls = refImages;
    }

    const response = await axios.post('https://api2img.shubiaobiao.com/v1/images/generations', payload, {
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${process.env.API_KEY}`
      },
      timeout: 120000
    });

    res.json(response.data);
  } catch (error) {
    console.error('Image Gen API Error:', error.response?.data || error.message);
    res.status(500).json({ error: 'Image generation failed' });
  }
});

// Endpoint to get image result
app.post('/api/get-image-result', async (req, res) => {
  try {
    const { id } = req.body;
    if (!id) return res.status(400).json({ error: 'Task ID is required' });

    const queryUrl = process.env.IMAGE_QUERY_API_URL || 'https://api.wuyinkeji.com/api/img/drawDetail';

    const response = await axios.get(queryUrl, {
      params: {
        key: process.env.API_KEY,
        id: id
      },
      headers: { 'Content-Type': 'application/json' }
    });

    res.json(response.data);
  } catch (error) {
    console.error('Image Query API Error:', error.response?.data || error.message);
    res.status(500).json({ error: 'Failed to query image result' });
  }
});

// ============================================
// 新增API端点 - SKU智能生图
// ============================================

/**
 * 获取Brain System Prompt
 */
app.get('/api/brain-prompt', (req, res) => {
  try {
    const promptPath = path.join(__dirname, 'prompts/brain-system.md');
    const content = fs.readFileSync(promptPath, 'utf-8');
    res.json({ content });
  } catch (error) {
    res.status(500).json({ error: 'Failed to read brain prompt' });
  }
});

/**
 * 获取示例SKU数据
 */
app.get('/api/sample-sku', (req, res) => {
  try {
    const skuPath = path.join(__dirname, 'data/sample-sku.json');
    const content = fs.readFileSync(skuPath, 'utf-8');
    res.json(JSON.parse(content));
  } catch (error) {
    res.status(500).json({ error: 'Failed to read sample SKU' });
  }
});

/**
 * Stage 1: Brain决策 - 生成图片规划
 * POST /api/plan-images
 * Body: { sku: SKU数据对象 }
 * Returns: { plan: Brain输出的图片规划 }
 */
app.post('/api/plan-images', async (req, res) => {
  try {
    const { sku } = req.body;
    
    if (!sku || !sku.sku_id) {
      return res.status(400).json({ error: 'SKU data is required' });
    }

    console.log(`Planning images for SKU: ${sku.sku_id}`);
    
    const plan = await planImages(sku, {
      apiKey: process.env.API_KEY
    });

    res.json({ 
      success: true,
      sku_id: sku.sku_id,
      plan 
    });
  } catch (error) {
    console.error('Plan Images Error:', error.message);
    res.status(500).json({ error: 'Failed to plan images: ' + error.message });
  }
});

/**
 * Stage 2: 生成单张图片（带约束注入）
 * POST /api/generate-image-with-constraints
 * Body: { sku, imageSpec, plan }
 * Returns: { imageUrl }
 */
app.post('/api/generate-image-with-constraints', async (req, res) => {
  try {
    const { sku, imageSpec } = req.body;
    
    if (!sku || !imageSpec) {
      return res.status(400).json({ error: 'SKU and imageSpec are required' });
    }

    // 注入约束
    const isComparison = imageSpec.is_comparison || imageSpec.id === 'APlus_02';
    const fullPrompt = injectConstraints(
      imageSpec.ai_prompt,
      sku,
      {
        isComparison,
        logoPlacement: imageSpec.logo_placement || {}
      }
    );

    // 确定宽高比
    const aspectRatio = getAspectRatio(imageSpec.id);

    console.log(`Generating image ${imageSpec.id} for SKU ${sku.sku_id}`);

    // 调用图像生成API
    const payload = {
      model: 'gemini-3-pro-image-preview',
      prompt: fullPrompt,
      n: 1,
      size: "1K",
      aspect_ratio: aspectRatio
    };

    // 添加参考图
    if (sku.ref_images && sku.ref_images.length > 0) {
      payload.image_urls = sku.ref_images;
    }

    const response = await axios.post('https://api2img.shubiaobiao.com/v1/images/generations', payload, {
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${process.env.API_KEY}`
      },
      timeout: 180000 // 3分钟超时
    });

    // 提取图片URL
    let imageUrl = null;
    if (response.data.data && Array.isArray(response.data.data) && response.data.data.length > 0) {
      imageUrl = response.data.data[0].url;
    }

    res.json({
      success: true,
      image_id: imageSpec.id,
      image_url: imageUrl,
      aspect_ratio: aspectRatio
    });

  } catch (error) {
    console.error('Generate Image With Constraints Error:', error.response?.data || error.message);
    res.status(500).json({ error: 'Image generation failed: ' + error.message });
  }
});

/**
 * 完整流程：为单个SKU生成全部12张图
 * POST /api/generate-sku
 * Body: { sku: SKU数据对象 }
 * Returns: { results: [{ id, url, status }] }
 */
app.post('/api/generate-sku', async (req, res) => {
  try {
    const { sku } = req.body;
    
    if (!sku || !sku.sku_id) {
      return res.status(400).json({ error: 'SKU data is required' });
    }

    console.log(`\n========== Starting SKU Generation: ${sku.sku_id} ==========`);

    // Stage 1: Brain规划
    console.log('Stage 1: Planning images...');
    const plan = await planImages(sku, { apiKey: process.env.API_KEY });
    console.log(`Plan generated with ${plan.images.length} images`);

    // Stage 2: 逐张生成图片
    console.log('Stage 2: Generating images...');
    const results = [];
    
    for (const imageSpec of plan.images) {
      try {
        console.log(`  Generating ${imageSpec.id}...`);
        
        // 注入约束
        const isComparison = imageSpec.is_comparison || imageSpec.id === 'APlus_02';
        const fullPrompt = injectConstraints(
          imageSpec.ai_prompt,
          sku,
          {
            isComparison,
            logoPlacement: imageSpec.logo_placement || {}
          }
        );

        const aspectRatio = getAspectRatio(imageSpec.id);

        const payload = {
          model: 'gemini-3-pro-image-preview',
          prompt: fullPrompt,
          n: 1,
          size: "1K",
          aspect_ratio: aspectRatio
        };

        if (sku.ref_images && sku.ref_images.length > 0) {
          payload.image_urls = sku.ref_images;
        }

        const response = await axios.post('https://api2img.shubiaobiao.com/v1/images/generations', payload, {
          headers: {
            'Content-Type': 'application/json',
            'Authorization': `Bearer ${process.env.API_KEY}`
          },
          timeout: 180000
        });

        let imageUrl = null;
        if (response.data.data && Array.isArray(response.data.data) && response.data.data.length > 0) {
          imageUrl = response.data.data[0].url;
        }

        results.push({
          id: imageSpec.id,
          type: imageSpec.type,
          purpose: imageSpec.purpose,
          url: imageUrl,
          status: imageUrl ? 'success' : 'failed',
          aspect_ratio: aspectRatio
        });

        console.log(`  ✓ ${imageSpec.id} generated`);

        // 添加小延迟避免API限流
        await new Promise(resolve => setTimeout(resolve, 1000));

      } catch (imgError) {
        console.error(`  ✗ ${imageSpec.id} failed:`, imgError.message);
        results.push({
          id: imageSpec.id,
          type: imageSpec.type,
          purpose: imageSpec.purpose,
          url: null,
          status: 'failed',
          error: imgError.message
        });
      }
    }

    const successCount = results.filter(r => r.status === 'success').length;
    console.log(`\n========== SKU Generation Complete: ${successCount}/${results.length} ==========\n`);

    res.json({
      success: true,
      sku_id: sku.sku_id,
      plan: {
        analysis: plan.analysis,
        image_count: plan.images.length
      },
      results,
      summary: {
        total: results.length,
        success: successCount,
        failed: results.length - successCount
      }
    });

  } catch (error) {
    console.error('Generate SKU Error:', error.message);
    res.status(500).json({ error: 'SKU generation failed: ' + error.message });
  }
});

/**
 * 批量生成：为多个SKU生成图片
 * POST /api/batch-generate
 * Body: { skus: [SKU数据数组], concurrency: 并发数 }
 * Returns: { results: [{ sku_id, status, images }] }
 */
app.post('/api/batch-generate', async (req, res) => {
  try {
    const { skus, concurrency = 1 } = req.body;
    
    if (!skus || !Array.isArray(skus) || skus.length === 0) {
      return res.status(400).json({ error: 'SKUs array is required' });
    }

    console.log(`\n========== Batch Generation Started: ${skus.length} SKUs ==========`);

    const allResults = [];
    
    // 简单的串行处理（避免API限流）
    for (let i = 0; i < skus.length; i++) {
      const sku = skus[i];
      console.log(`\nProcessing SKU ${i + 1}/${skus.length}: ${sku.sku_id}`);
      
      try {
        // 调用单个SKU生成
        const plan = await planImages(sku, { apiKey: process.env.API_KEY });
        
        const skuResults = {
          sku_id: sku.sku_id,
          status: 'processing',
          images: [],
          plan_analysis: plan.analysis
        };

        for (const imageSpec of plan.images) {
          try {
            const isComparison = imageSpec.is_comparison || imageSpec.id === 'APlus_02';
            const fullPrompt = injectConstraints(
              imageSpec.ai_prompt,
              sku,
              { isComparison, logoPlacement: imageSpec.logo_placement || {} }
            );

            const aspectRatio = getAspectRatio(imageSpec.id);

            const payload = {
              model: 'gemini-3-pro-image-preview',
              prompt: fullPrompt,
              n: 1,
              size: "1K",
              aspect_ratio: aspectRatio
            };

            if (sku.ref_images && sku.ref_images.length > 0) {
              payload.image_urls = sku.ref_images;
            }

            const response = await axios.post('https://api2img.shubiaobiao.com/v1/images/generations', payload, {
              headers: {
                'Content-Type': 'application/json',
                'Authorization': `Bearer ${process.env.API_KEY}`
              },
              timeout: 180000
            });

            let imageUrl = null;
            if (response.data.data && Array.isArray(response.data.data) && response.data.data.length > 0) {
              imageUrl = response.data.data[0].url;
            }

            skuResults.images.push({
              id: imageSpec.id,
              url: imageUrl,
              status: imageUrl ? 'success' : 'failed'
            });

            // 延迟避免限流
            await new Promise(resolve => setTimeout(resolve, 1500));

          } catch (imgErr) {
            skuResults.images.push({
              id: imageSpec.id,
              url: null,
              status: 'failed',
              error: imgErr.message
            });
          }
        }

        const successCount = skuResults.images.filter(img => img.status === 'success').length;
        skuResults.status = successCount === skuResults.images.length ? 'complete' : 'partial';
        skuResults.summary = {
          total: skuResults.images.length,
          success: successCount,
          failed: skuResults.images.length - successCount
        };

        allResults.push(skuResults);
        console.log(`  SKU ${sku.sku_id} complete: ${successCount}/${skuResults.images.length} images`);

      } catch (skuErr) {
        console.error(`  SKU ${sku.sku_id} failed:`, skuErr.message);
        allResults.push({
          sku_id: sku.sku_id,
          status: 'failed',
          error: skuErr.message,
          images: []
        });
      }
    }

    const completedCount = allResults.filter(r => r.status === 'complete').length;
    console.log(`\n========== Batch Complete: ${completedCount}/${skus.length} SKUs ==========\n`);

    res.json({
      success: true,
      total_skus: skus.length,
      completed: completedCount,
      results: allResults
    });

  } catch (error) {
    console.error('Batch Generate Error:', error.message);
    res.status(500).json({ error: 'Batch generation failed: ' + error.message });
  }
});

/**
 * 获取约束预览（调试用）
 * POST /api/preview-constraints
 */
app.post('/api/preview-constraints', (req, res) => {
  try {
    const { sku, imageSpec } = req.body;
    
    if (!sku) {
      return res.status(400).json({ error: 'SKU data is required' });
    }

    const isComparison = imageSpec?.is_comparison || imageSpec?.id === 'APlus_02';
    const fullPrompt = injectConstraints(
      imageSpec?.ai_prompt || 'Sample prompt for preview',
      sku,
      {
        isComparison,
        logoPlacement: imageSpec?.logo_placement || {}
      }
    );

    res.json({
      original_prompt: imageSpec?.ai_prompt || 'Sample prompt for preview',
      full_prompt_with_constraints: fullPrompt,
      prompt_length: fullPrompt.length
    });
  } catch (error) {
    res.status(500).json({ error: error.message });
  }
});

// ============================================
// 启动服务器
// ============================================

app.listen(port, () => {
  console.log(`\n🚀 Server running at http://localhost:${port}`);
  console.log('📁 Available endpoints:');
  console.log('   GET  /api/prompt-template     - Get prompt template');
  console.log('   GET  /api/brain-prompt        - Get brain system prompt');
  console.log('   GET  /api/sample-sku          - Get sample SKU data');
  console.log('   POST /api/plan-images         - Stage 1: Brain planning');
  console.log('   POST /api/generate-sku        - Full SKU generation (12 images)');
  console.log('   POST /api/batch-generate      - Batch SKU generation');
  console.log('   POST /api/preview-constraints - Preview constraints injection');
  console.log('');
});