/**
 * POC Workflow V4: 完整优化版
 * 
 * 改进：
 * 1. Vision超时重试机制
 * 2. 可配置的套路模板系统
 * 3. 基于真实交付成品的Prompt
 * 4. 强化产品一致性约束
 * 
 * 用法：
 * node poc-workflow-v4.js --material-dir="./素材/素材/已有的素材" --sku-name="Touchdog冰蓝色伊丽莎白圈"
 */

require('dotenv').config();
const axios = require('axios');
const fs = require('fs');
const path = require('path');
const { S3Client, PutObjectCommand } = require('@aws-sdk/client-s3');
const { extractProductFeatures, buildGoldenDescription } = require('./lib/vision-extractor');
const { generateAllPrompts } = require('./lib/template-config');

// ============================================================
// 配置
// ============================================================
const API_KEY = 'G9rXx3Ag2Xfa7Gs8zou6t6HqeZ';
const API_BASE = 'https://api.wuyinkeji.com/api';

// R2客户端
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,
});

// ============================================================
// 命令行参数解析
// ============================================================
function parseArgs() {
  const args = process.argv.slice(2);
  const config = {
    materialDir: path.join(__dirname, '素材/素材/已有的素材'),
    skuName: 'Touchdog 软质伊丽莎白圈',
    brandName: 'TOUCHDOG',
    productName: 'CAT SOFT CONE COLLAR',
    petType: 'cat',
    outputDir: null
  };

  for (const arg of args) {
    if (arg.startsWith('--material-dir=')) {
      config.materialDir = arg.split('=')[1];
    } else if (arg.startsWith('--sku-name=')) {
      config.skuName = arg.split('=')[1];
    } else if (arg.startsWith('--brand-name=')) {
      config.brandName = arg.split('=')[1];
    } else if (arg.startsWith('--product-name=')) {
      config.productName = arg.split('=')[1];
    } else if (arg.startsWith('--output-dir=')) {
      config.outputDir = arg.split('=')[1];
    } else if (arg.startsWith('--pet-type=')) {
      config.petType = arg.split('=')[1];
    }
  }

  // 自动生成输出目录
  if (!config.outputDir) {
    const timestamp = new Date().toISOString().slice(0, 19).replace(/[T:]/g, '-');
    const safeName = config.skuName.replace(/[^a-zA-Z0-9\u4e00-\u9fa5]/g, '_').slice(0, 20);
    config.outputDir = path.join(__dirname, `output_v4_${safeName}_${timestamp}`);
  }

  return config;
}

// ============================================================
// 工具函数
// ============================================================

async function uploadToR2(filePath) {
  const fileName = `v4-${Date.now()}-${path.basename(filePath)}`;
  const fileBuffer = fs.readFileSync(filePath);
  
  await r2Client.send(new PutObjectCommand({
    Bucket: process.env.R2_BUCKET_NAME || 'ai-flow',
    Key: fileName,
    Body: fileBuffer,
    ContentType: filePath.endsWith('.png') ? 'image/png' : 'image/jpeg',
  }));

  return process.env.R2_PUBLIC_DOMAIN
    ? `${process.env.R2_PUBLIC_DOMAIN}/${fileName}`
    : `https://${process.env.R2_ACCOUNT_ID}.r2.cloudflarestorage.com/${process.env.R2_BUCKET_NAME}/${fileName}`;
}

// 扫描素材目录
async function scanMaterials(materialDir) {
  console.log('📁 扫描素材目录:', materialDir);
  
  if (!fs.existsSync(materialDir)) {
    throw new Error(`素材目录不存在: ${materialDir}`);
  }
  
  const files = fs.readdirSync(materialDir);
  const images = files.filter(f => /\.(jpg|jpeg|png)$/i.test(f));
  
  console.log(`   找到 ${images.length} 张图片: ${images.join(', ')}`);
  
  // 识别图片类型
  let flatImage = null;  // 平铺图
  let wornImage = null;  // 佩戴图
  
  for (const img of images) {
    const imgPath = path.join(materialDir, img);
    const imgLower = img.toLowerCase();
    
    // 平铺图识别
    if (img.includes('5683') || imgLower.includes('flat') || imgLower.includes('1.png')) {
      flatImage = imgPath;
    }
    // 佩戴图识别
    if (img.includes('6514') || imgLower.includes('worn') || imgLower.includes('wear') || imgLower.includes('3.png')) {
      wornImage = imgPath;
    }
  }
  
  // 如果没找到，用启发式规则
  if (!flatImage) {
    // 优先选png格式（通常更清晰）
    const pngImages = images.filter(i => i.toLowerCase().endsWith('.png'));
    flatImage = pngImages.length > 0 
      ? path.join(materialDir, pngImages[0])
      : path.join(materialDir, images[0]);
  }
  
  if (!wornImage) {
    const jpgImages = images.filter(i => i.toLowerCase().endsWith('.jpg') || i.toLowerCase().endsWith('.jpeg'));
    wornImage = jpgImages.length > 0
      ? path.join(materialDir, jpgImages[0])
      : flatImage;
  }
  
  console.log('   ✓ 平铺图:', path.basename(flatImage));
  console.log('   ✓ 佩戴图:', path.basename(wornImage));
  
  return { flatImage, wornImage, allImages: images.map(i => path.join(materialDir, i)) };
}

// 生图任务
async function submitImageTask(prompt, refImageUrl, aspectRatio = '1:1') {
  const payload = {
    key: API_KEY,
    prompt: prompt,
    img_url: refImageUrl,
    aspectRatio: aspectRatio,
    imageSize: '1K'
  };

  const response = await axios.post(`${API_BASE}/img/nanoBanana-pro`, payload, {
    timeout: 30000
  });
  
  const taskId = response.data.data?.id;
  if (!taskId) {
    console.error('Submit response:', response.data);
    throw new Error('No task ID returned');
  }
  return taskId;
}

async function pollImageResult(taskId) {
  let attempts = 0;
  const maxAttempts = 90;
  
  while (attempts < maxAttempts) {
    try {
      const response = await axios.get(`${API_BASE}/img/drawDetail`, {
        params: { key: API_KEY, id: taskId },
        timeout: 10000
      });
      
      const data = response.data.data;
      
      if (data && data.status === 2 && data.image_url) {
        return data.image_url;
      } else if (data && data.status === 3) {
        throw new Error('Generation failed: ' + (data.fail_reason || 'Unknown'));
      }
      
      process.stdout.write('.');
      await new Promise(r => setTimeout(r, 2000));
      attempts++;
    } catch (error) {
      if (error.message.includes('Generation failed')) throw error;
      // 网络错误，继续重试
      process.stdout.write('x');
      await new Promise(r => setTimeout(r, 3000));
      attempts++;
    }
  }
  throw new Error('Timeout after ' + maxAttempts + ' attempts');
}

async function generateSingleImage(item, refUrl, outputDir) {
  console.log(`\n🎨 [${item.id}] ${item.name}`);
  
  let retryCount = 0;
  const maxRetries = 2;
  
  while (retryCount <= maxRetries) {
    try {
      if (retryCount > 0) console.log(`   重试 ${retryCount}/${maxRetries}...`);
      
      console.log('   提交任务...');
      const taskId = await submitImageTask(item.prompt, refUrl, item.aspectRatio);
      console.log(`   Task ID: ${taskId}`);
      
      const imageUrl = await pollImageResult(taskId);
      console.log('\n   ✓ 生成成功');
      
      // 下载保存
      const imgRes = await axios.get(imageUrl, { responseType: 'arraybuffer', timeout: 30000 });
      const outputPath = path.join(outputDir, `${item.id}.jpg`);
      fs.writeFileSync(outputPath, imgRes.data);
      console.log(`   ✓ 保存: ${item.id}.jpg`);
      
      return { id: item.id, success: true, path: outputPath };
      
    } catch (error) {
      console.error(`\n   ✗ 失败: ${error.message}`);
      retryCount++;
      if (retryCount <= maxRetries) {
        await new Promise(r => setTimeout(r, 5000));
      }
    }
  }
  
  return { id: item.id, success: false, error: 'Max retries exceeded' };
}

// ============================================================
// 主流程
// ============================================================

async function main() {
  const config = parseArgs();
  
  console.log('\n' + '═'.repeat(70));
  console.log('🚀 POC Workflow V4 - 完整优化版');
  console.log('═'.repeat(70));
  console.log('\n📋 配置:');
  console.log('   素材目录:', config.materialDir);
  console.log('   SKU名称:', config.skuName);
  console.log('   品牌:', config.brandName);
  console.log('   输出目录:', config.outputDir);
  
  // 创建输出目录
  if (!fs.existsSync(config.outputDir)) {
    fs.mkdirSync(config.outputDir, { recursive: true });
  }

  // ========== 阶段1: 扫描素材 ==========
  console.log('\n' + '─'.repeat(70));
  console.log('📁 阶段1: 扫描素材');
  console.log('─'.repeat(70));
  
  const materials = await scanMaterials(config.materialDir);
  
  if (!materials.flatImage) {
    console.error('❌ 未找到素材图片');
    process.exit(1);
  }

  // ========== 阶段2: 上传素材 ==========
  console.log('\n' + '─'.repeat(70));
  console.log('📤 阶段2: 上传素材到云存储');
  console.log('─'.repeat(70));
  
  const flatImageUrl = await uploadToR2(materials.flatImage);
  console.log('   平铺图:', flatImageUrl);
  
  let wornImageUrl = flatImageUrl;
  if (materials.wornImage && materials.wornImage !== materials.flatImage) {
    wornImageUrl = await uploadToR2(materials.wornImage);
    console.log('   佩戴图:', wornImageUrl);
  }

  // ========== 阶段3: Vision分析 ==========
  console.log('\n' + '─'.repeat(70));
  console.log('🔍 阶段3: Vision分析产品特征 (超时重试机制)');
  console.log('─'.repeat(70));
  
  const visionResult = await extractProductFeatures(flatImageUrl, {
    maxRetries: 3,
    timeout: 150000,  // 150秒
    retryDelay: 8000,
    cacheDir: config.outputDir,
    cacheKey: path.basename(materials.flatImage).replace(/\.[^.]+$/, '')
  });
  
  const goldenDescription = buildGoldenDescription(visionResult);
  
  console.log('\n📝 Golden Description预览:');
  console.log('─'.repeat(50));
  console.log(goldenDescription.substring(0, 600) + '...');
  
  // 保存分析结果
  fs.writeFileSync(
    path.join(config.outputDir, 'vision-analysis.json'),
    JSON.stringify(visionResult, null, 2)
  );
  fs.writeFileSync(
    path.join(config.outputDir, 'golden-description.txt'),
    goldenDescription
  );

  // ========== 阶段4: 生成Prompts ==========
  console.log('\n' + '─'.repeat(70));
  console.log('📝 阶段4: 生成12张图的Prompts (可配置模板系统)');
  console.log('─'.repeat(70));
  
  const product = { goldenDescription };
  const skuInfo = {
    brandName: config.brandName,
    productName: config.productName,
    petType: config.petType
  };
  
  const prompts = generateAllPrompts(product, skuInfo);
  console.log(`   ✓ 生成了 ${prompts.length} 个Prompt`);
  
  // 保存prompts供参考
  fs.writeFileSync(
    path.join(config.outputDir, 'all-prompts.json'),
    JSON.stringify(prompts.map(p => ({
      id: p.id,
      name: p.name,
      type: p.type,
      aspectRatio: p.aspectRatio,
      prompt: p.prompt
    })), null, 2)
  );

  // ========== 阶段5: 批量生成 ==========
  console.log('\n' + '─'.repeat(70));
  console.log('🎨 阶段5: 批量生成图片');
  console.log('─'.repeat(70));
  
  const results = [];
  for (const promptItem of prompts) {
    // Main_02用平铺图参考，其他用佩戴图
    const refUrl = promptItem.id === 'Main_02' ? flatImageUrl : wornImageUrl;
    
    const result = await generateSingleImage(promptItem, refUrl, config.outputDir);
    results.push(result);
  }

  // ========== 阶段6: 总结 ==========
  console.log('\n' + '═'.repeat(70));
  console.log('📊 生成完成');
  console.log('═'.repeat(70));
  
  const successCount = results.filter(r => r.success).length;
  console.log(`\n✅ 成功: ${successCount}/${results.length}`);
  
  if (successCount < results.length) {
    console.log('\n❌ 失败的图片:');
    results.filter(r => !r.success).forEach(r => {
      console.log(`   - ${r.id}: ${r.error}`);
    });
  }
  
  console.log(`\n📁 输出目录: ${config.outputDir}`);
  
  // 保存结果摘要
  fs.writeFileSync(
    path.join(config.outputDir, 'results-summary.json'),
    JSON.stringify({
      config,
      visionSuccess: !!visionResult,
      results,
      summary: {
        total: results.length,
        success: successCount,
        failed: results.length - successCount
      },
      timestamp: new Date().toISOString()
    }, null, 2)
  );
  
  console.log('\n💡 提示: 查看 all-prompts.json 了解每张图的Prompt详情');
}

main().catch(console.error);