/**
 * POC Workflow V3: 通用化工作流
 * 
 * 特性：
 * 1. 素材目录作为参数，支持一键切换SKU
 * 2. Vision自动提取产品特征
 * 3. 基于真实交付成品的12张图模板
 * 4. 带文字排版的A+/主图
 * 
 * 用法：
 * node poc-workflow-v3.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 { generateAllPrompts } = require('./prompts/image-templates');

// ============================================================
// 配置
// ============================================================
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',
    outputDir: null, // 自动生成
    sellingPoints: [
      'CLOUD-LIGHT COMFORT',
      'WIDER & CLEARER',
      'FOLDABLE & PORTABLE',
      'WATERPROOF & EASY CLEAN'
    ],
    competitorWeaknesses: [
      'HEAVY & BULKY',
      'BLOCKS VISION & MOVEMENT',
      'HARD TO STORE'
    ],
    features: [
      { title: 'STURDY AND BREATHABLE', desc: ', DURABLE AND COMFORTABLE' },
      { title: 'EASY TO CLEAN, STYLISH', desc: 'AND ATTRACTIVE' },
      { title: 'REINFORCED STITCHING PROCESS', desc: 'AND DURABLE FABRIC' }
    ],
    sizeChart: {
      XS: { neck: '5.6-6.8IN', depth: '3.2IN' },
      S: { neck: '7.2-8.4IN', depth: '4IN' },
      M: { neck: '8.8-10.4IN', depth: '5IN' },
      L: { neck: '10.8-12.4IN', depth: '6IN' },
      XL: { neck: '12.8-14.4IN', depth: '7IN' }
    }
  };

  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];
    }
  }

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

  return config;
}

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

async function uploadToR2(filePath) {
  const fileName = `v3-${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);
  
  const files = fs.readdirSync(materialDir);
  const images = files.filter(f => /\.(jpg|jpeg|png)$/i.test(f));
  
  console.log(`   找到 ${images.length} 张图片`);
  
  // 识别图片类型
  let flatImage = null;  // 平铺图
  let wornImage = null;  // 佩戴图
  
  for (const img of images) {
    const imgPath = path.join(materialDir, img);
    // 简单启发式：文件名包含特定关键词
    if (img.includes('5683') || img.toLowerCase().includes('flat')) {
      flatImage = imgPath;
    } else if (img.includes('6514') || img.toLowerCase().includes('worn') || img.toLowerCase().includes('wear')) {
      wornImage = imgPath;
    }
  }
  
  // 如果没找到，用第一张和第二张
  if (!flatImage && images.length > 0) {
    flatImage = path.join(materialDir, images[0]);
  }
  if (!wornImage && images.length > 1) {
    wornImage = path.join(materialDir, images[1]);
  } else if (!wornImage) {
    wornImage = flatImage;
  }
  
  console.log('   平铺图:', flatImage ? path.basename(flatImage) : '未找到');
  console.log('   佩戴图:', wornImage ? path.basename(wornImage) : '未找到');
  
  return { flatImage, wornImage, allImages: images.map(i => path.join(materialDir, i)) };
}

// Vision API提取产品特征
const VISION_EXTRACT_PROMPT = `Analyze this pet recovery cone product image in EXTREME detail.

Output ONLY a valid JSON object (no markdown, no explanation):

{
  "color": {
    "primary": "<hex code>",
    "name": "<descriptive name>",
    "secondary": "<accent colors>"
  },
  "shape": {
    "type": "<flower/fan/cone>",
    "petal_count": <number>,
    "opening": "<C-shaped/full-circle>",
    "description": "<detailed shape>"
  },
  "material": {
    "type": "<PU/fabric/plastic>",
    "finish": "<glossy/matte/satin>",
    "texture": "<smooth/quilted>"
  },
  "edge_binding": {
    "color": "<color>",
    "material": "<ribbed/fabric>"
  },
  "closure": {
    "type": "<velcro/button>",
    "color": "<color>",
    "position": "<location>"
  },
  "logo": {
    "text": "<brand>",
    "style": "<embroidered/printed>",
    "position": "<location>"
  },
  "unique_features": ["<list>"],
  "overall_description": "<2-3 sentence summary>"
}`;

async function extractProductFeatures(imageUrl) {
  console.log('\n🔍 Vision分析产品特征...');
  
  try {
    const response = await axios.post(`${API_BASE}/chat/index`, {
      key: API_KEY,
      model: 'gemini-3-pro',
      content: VISION_EXTRACT_PROMPT,
      image_url: imageUrl
    }, { timeout: 90000 });

    const content = response.data.data?.choices?.[0]?.message?.content || 
                    response.data.data?.content;
    
    if (!content) {
      throw new Error('Vision响应为空');
    }

    // 提取JSON
    const jsonMatch = content.match(/\{[\s\S]*\}/);
    if (jsonMatch) {
      const parsed = JSON.parse(jsonMatch[0]);
      console.log('   ✓ 提取成功');
      return parsed;
    }
    
    throw new Error('无法解析JSON');
  } catch (error) {
    console.error('   ✗ Vision分析失败:', error.message);
    return null;
  }
}

// 将Vision结果转换为Golden Description
function buildGoldenDescription(visionResult) {
  if (!visionResult) {
    return `
PRODUCT: Pet recovery cone collar
- Soft, flexible material
- C-shaped opening design
- Velcro closure
- Comfortable for pets
`;
  }

  const r = visionResult;
  return `
EXACT PRODUCT APPEARANCE (AUTO-EXTRACTED):
- Shape: ${r.shape?.petal_count || '?'}-PETAL ${r.shape?.type?.toUpperCase() || 'FLOWER'} shape, ${r.shape?.opening || 'C-shaped'} opening
- Color: ${r.color?.name?.toUpperCase() || 'UNKNOWN'} (${r.color?.primary || '#???'})
- Material: ${r.material?.finish || 'Soft'} ${r.material?.type || 'fabric'} with ${r.material?.texture || 'smooth'} texture
- Edge binding: ${r.edge_binding?.color || 'Contrasting color'} ${r.edge_binding?.material || 'ribbed elastic'} around inner neck hole
- Closure: ${r.closure?.color || 'White'} ${r.closure?.type || 'velcro'} on ${r.closure?.position || 'one end'}
- Logo: "${r.logo?.text || 'TOUCHDOG'}" ${r.logo?.style || 'embroidered'}

UNIQUE FEATURES:
${(r.unique_features || []).map(f => `- ${f}`).join('\n') || '- Soft, comfortable design'}

CRITICAL PROHIBITIONS:
- ❌ NO printed patterns or colorful fabric designs  
- ❌ NO hard plastic transparent cones
- ❌ NO fully circular/closed shapes (must have C-opening)
- ❌ NO random brand logos
`.trim();
}

// 生图任务
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);
  const taskId = response.data.data?.id;
  
  if (!taskId) {
    throw new Error('No task ID returned');
  }
  return taskId;
}

async function pollImageResult(taskId) {
  let attempts = 0;
  const maxAttempts = 90;
  
  while (attempts < maxAttempts) {
    const response = await axios.get(`${API_BASE}/img/drawDetail`, {
      params: { key: API_KEY, id: taskId }
    });
    
    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++;
  }
  throw new Error('Timeout');
}

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' });
      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++;
      await new Promise(r => setTimeout(r, 3000));
    }
  }
  
  return { id: item.id, success: false, error: 'Max retries exceeded' };
}

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

async function main() {
  const config = parseArgs();
  
  console.log('='.repeat(70));
  console.log('🚀 POC Workflow V3 - 通用化工作流');
  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. 上传素材到R2
  console.log('\n' + '─'.repeat(70));
  console.log('📤 阶段2: 上传素材');
  console.log('─'.repeat(70));
  
  const flatImageUrl = await uploadToR2(materials.flatImage);
  console.log('   平铺图URL:', flatImageUrl);
  
  let wornImageUrl = flatImageUrl;
  if (materials.wornImage && materials.wornImage !== materials.flatImage) {
    wornImageUrl = await uploadToR2(materials.wornImage);
    console.log('   佩戴图URL:', wornImageUrl);
  }

  // 3. Vision分析
  console.log('\n' + '─'.repeat(70));
  console.log('🔍 阶段3: Vision分析产品特征');
  console.log('─'.repeat(70));
  
  const visionResult = await extractProductFeatures(flatImageUrl);
  const goldenDescription = buildGoldenDescription(visionResult);
  
  console.log('\n生成的Golden Description:');
  console.log(goldenDescription.substring(0, 500) + '...');
  
  // 保存Vision结果
  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. 生成12张图的Prompts
  console.log('\n' + '─'.repeat(70));
  console.log('📝 阶段4: 生成Prompts');
  console.log('─'.repeat(70));
  
  const product = { goldenDescription };
  const skuInfo = {
    brandName: config.brandName,
    productName: config.productName,
    sellingPoints: config.sellingPoints,
    competitorWeaknesses: config.competitorWeaknesses,
    features: config.features,
    sizeChart: config.sizeChart
  };
  
  const prompts = generateAllPrompts(product, skuInfo);
  console.log(`   生成了 ${prompts.length} 个Prompt`);
  
  // 保存prompts供参考
  fs.writeFileSync(
    path.join(config.outputDir, 'prompts.json'),
    JSON.stringify(prompts.map(p => ({ id: p.id, name: p.name, aspectRatio: p.aspectRatio })), null, 2)
  );

  // 5. 批量生成图片
  console.log('\n' + '─'.repeat(70));
  console.log('🎨 阶段5: 批量生成图片');
  console.log('─'.repeat(70));
  
  const results = [];
  for (const promptItem of prompts) {
    // 根据图片类型选择参考图
    const refUrl = promptItem.id.includes('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.json'),
    JSON.stringify({
      config,
      visionResult,
      results,
      summary: {
        total: results.length,
        success: successCount,
        failed: results.length - successCount
      }
    }, null, 2)
  );
}

main().catch(console.error);