/**
 * POC Workflow: Vision Analysis -> Optimized Generation (Full 12 Images)
 * API Provider: wuyinkeji.com
 */

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

// 配置
const API_KEY = 'G9rXx3Ag2Xfa7Gs8zou6t6HqeZ';
const API_BASE = 'https://api.wuyinkeji.com/api';
const OUTPUT_DIR = path.join(__dirname, 'output_poc_full');
const MATERIAL_DIR = path.join(__dirname, '素材/素材/已有的素材');

// 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,
});

async function uploadToR2(filePath) {
  const fileName = `poc-${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',
  }));

  const publicUrl = 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}`;
  
  return publicUrl;
}

// 核心函数1：视觉分析 (Vision Agent)
async function analyzeProductVisuals(refImageUrls) {
  console.log('👁️ 正在进行视觉分析 (Vision Agent)...');
  
  const analysisPrompt = `
You are a professional e-commerce product photographer.
Analyze these reference images of a pet recovery cone collar.
Extract visual characteristics:
1. Material & Texture: (e.g. glossy PU, matte cotton)
2. Structure (Flat): (e.g. Open C-shape, Fan shape)
3. Structure (Worn): (e.g. Cone shape flared OUTWARDS/FORWARD)
4. Key Details: (e.g. logo, binding color)

Output JSON: { "material": "...", "structure_flat": "...", "structure_worn": "...", "details": "..." }
`;

  try {
    // 构造Wuyin Vision请求
    // 注意：根据文档截图，/api/chat/index 支持 image_url 参数
    // 这里我们只传第一张图作为主要参考，或者尝试传多张如果API支持
    // 由于接口文档通常 image_url 是字符串，我们先传最具代表性的图（通常是第一张）
    // 如果需要多图分析，可能需要多次调用或拼接
    
    // 这里简单起见，我们取一张最清晰的图（假设是第一张或摊平图）进行分析
    const mainRefImage = refImageUrls.find(url => url.includes('5683')) || refImageUrls[0];

    const response = await axios.post(
      `${API_BASE}/chat/index`,
      {
        key: API_KEY,
        model: 'gemini-3-pro', // 明确指定 vision 模型为 gemini-3-pro
        content: analysisPrompt,
        image_url: mainRefImage
      },
      { timeout: 60000 }
    );

    let content = response.data.data?.content || response.data.choices?.[0]?.message?.content;
    
    if (!content) throw new Error('No content in vision response');

    content = content.replace(/```json|```/g, '').trim();
    
    // 尝试解析JSON，如果失败则用正则提取或兜底
    let analysis;
    try {
        analysis = JSON.parse(content);
    } catch (e) {
        console.warn('JSON parse failed, using raw content');
        analysis = {
            material: "Smooth waterproof PU, slightly glossy",
            structure_flat: "Open C-shape/Fan shape",
            structure_worn: "Cone flared upwards/outwards",
            details: "Embroidered Touchdog logo"
        };
    }
    
    console.log('✅ 视觉分析完成:', analysis);
    return analysis;

  } catch (error) {
    console.error('❌ 视觉分析失败:', error.message);
    // 兜底描述
    return {
      material: "Smooth waterproof PU fabric with a slight sheen",
      structure_flat: "Open C-shape (Fan-like) with petal segments",
      structure_worn: "Cone shape flared UPWARDS/FORWARD around head",
      details: "Embroidered Touchdog logo on right petal"
    };
  }
}

// 生图任务提交
async function submitImageTask(prompt, refImageUrls, aspectRatio = '1:1') {
  try {
    // Wuyin API: /api/img/nanoBanana-pro
    // 参数: key, prompt, img_url (参考图), aspectRatio
    
    // 注意：img_url 只能传一张。我们传最相关的一张。
    // 对于场景图，传佩戴图；对于平铺图，传平铺图。
    let refImg = refImageUrls[0];
    if (prompt.includes('flat') && refImageUrls.some(u => u.includes('5683'))) {
        refImg = refImageUrls.find(u => u.includes('5683')); // 摊平图
    } else if (refImageUrls.some(u => u.includes('6514'))) {
        refImg = refImageUrls.find(u => u.includes('6514')); // 佩戴图
    }

    const payload = {
      key: API_KEY,
      prompt: prompt,
      img_url: refImg, // 传单张参考图强约束
      aspectRatio: aspectRatio,
      imageSize: '1K'
    };

    console.log(`   提交任务... (Ref: ${path.basename(refImg || '')})`);
    
    const response = await axios.post(`${API_BASE}/img/nanoBanana-pro`, payload);
    
    // 假设返回格式 { code: 200, data: { id: "..." }, msg: "success" }
    // 或者直接返回 task_id，需根据实际响应调整
    // 根据截图: 异步任务通常返回 id
    
    // 这里假设 API 返回包含 task id
    const taskId = response.data.data?.id; // 修正: 获取 data.id
    
    if (!taskId) {
        console.error('Submit response:', response.data);
        throw new Error('No task ID returned');
    }
    
    return taskId;

  } catch (error) {
    console.error('Submit failed:', error.message);
    throw error;
  }
}

// 轮询图片结果
async function pollImageResult(taskId) {
  let attempts = 0;
  const maxAttempts = 60; // 60 * 2s = 2 mins
  
  while (attempts < maxAttempts) {
    try {
      const response = await axios.get(`${API_BASE}/img/drawDetail`, {
        params: { key: API_KEY, id: taskId }
      });
      
      // 状态判断修正: status 2 是成功
      const data = response.data.data;
      
      if (data && data.status === 2 && data.image_url) {
        return data.image_url;
      } else if (data && (data.status === 3 || data.status === 'failed')) { // 假设 3 是失败，或者保留 failed 字符串判断
        throw new Error('Generation failed: ' + (data.fail_reason || 'Unknown'));
      }
      
      console.log(`   状态: ${data.status} (等待中...)`);
      
      // Still processing
      await new Promise(r => setTimeout(r, 2000));
      attempts++;
      process.stdout.write('.');
      
    } catch (error) {
      console.error('Poll error:', error.message);
      throw error; // 直接抛出，触发重试
    }
  }
  throw new Error('Timeout waiting for image');
}

// 生成单张图（含重试）
async function generateSingleImage(item, refImageUrls) {
  console.log(`\n🎨 生成 ${item.id} (${item.type})...`);
  
  let retryCount = 0;
  const maxRetries = 2;
  
  while (retryCount <= maxRetries) {
    try {
      if (retryCount > 0) console.log(`   重试第 ${retryCount} 次...`);
      
      const taskId = await submitImageTask(item.prompt, refImageUrls, item.aspectRatio);
      console.log(`   Task ID: ${taskId}`);
      
      const imageUrl = await pollImageResult(taskId);
      console.log('   ✓ 生成成功');
      
      // 下载保存
      const imgRes = await axios.get(imageUrl, { responseType: 'arraybuffer' });
      fs.writeFileSync(path.join(OUTPUT_DIR, `${item.id}.jpg`), imgRes.data);
      console.log('   ✓ 保存本地');
      return; // 成功退出
      
    } catch (error) {
      console.error(`   ✗ 失败: ${error.message}`);
      retryCount++;
      await new Promise(r => setTimeout(r, 3000));
    }
  }
  console.error(`   ❌ ${item.id} 最终失败`);
}

// 主流程
async function main() {
  if (!fs.existsSync(OUTPUT_DIR)) fs.mkdirSync(OUTPUT_DIR, { recursive: true });

  // 1. 上传
  console.log('📤 Step 1: 上传素材...');
  const files = fs.readdirSync(MATERIAL_DIR).filter(f => /\.(jpg|png)$/i.test(f));
  const refImageUrls = [];
  files.sort(); 
  
  for (const file of files) {
    if (file.startsWith('.')) continue;
    const url = await uploadToR2(path.join(MATERIAL_DIR, file));
    refImageUrls.push(url);
  }
  console.log(`   共 ${refImageUrls.length} 张`);

  // 2. 视觉分析
  const va = await analyzeProductVisuals(refImageUrls);
  // const va = {
  //     material: "Smooth waterproof PU fabric with a slight glossy sheen",
  //     structure_flat: "Open C-shape (Fan-like) with petal segments, NOT a closed circle",
  //     structure_worn: "Cone shape flared UPWARDS/FORWARD around head (megaphone shape)",
  //     details: "Embroidered Touchdog logo on right petal, teal edge binding"
  // };

  // 3. 定义12张图的Prompt (结合视觉分析)
  const prompts = [
    // --- 主图 (1:1) ---
    {
      id: 'Main_01_Hero',
      type: '场景首图',
      aspectRatio: '1:1',
      prompt: `Amazon main image. Ragdoll cat wearing ${va.material} cone collar. 
      CRITICAL: Cone opening must face FORWARD/OUTWARDS around the cat's face (like a megaphone), NOT wrapped like a scarf. 
      Cat is looking comfortable. Warm cozy home background. 
      Product texture: ${va.material}. Structure: ${va.structure_worn}. 
      Logo: Touchdog embroidered on right side. High quality, 8k.`
    },
    {
      id: 'Main_02_Flat',
      type: '平铺图',
      aspectRatio: '1:1',
      prompt: `Product flat lay, white background. 
      Structure: ${va.structure_flat} (Open C-shape/Fan-like, NOT closed circle).
      Texture: ${va.material}, glossy PU sheen.
      Details: Velcro strap visible, ${va.details}. Clean studio light.`
    },
    {
      id: 'Main_03_Function',
      type: '功能演示',
      aspectRatio: '1:1',
      prompt: `Product feature shot. Close-up of the velcro strap adjustment. 
      Hands adjusting the strap for secure fit. 
      Show the ${va.material} texture and stitching quality. 
      Light blue background with text overlay "ADJUSTABLE FIT" (optional).`
    },
    {
      id: 'Main_04_Scenarios',
      type: '场景拼图',
      aspectRatio: '1:1',
      prompt: `Split screen composition. 
      Top: Cat eating food easily while wearing the cone (cone facing forward).
      Bottom: Cat sleeping comfortably. 
      Show flexibility and comfort. Soft lighting.`
    },
    {
      id: 'Main_05_Size',
      type: '尺寸图',
      aspectRatio: '1:1',
      prompt: `Product infographic. The ${va.structure_flat} cone laid flat. 
      Ruler graphics measuring the depth (24.5cm). 
      Size chart table on the side: XS, S, M, L, XL. 
      Professional clean design.`
    },
    {
      id: 'Main_06_Brand',
      type: '品牌汇总',
      aspectRatio: '1:1',
      prompt: `Brand image. Large Touchdog logo in center. 
      Product floating in background. 
      Icons: Feather (Lightweight), Water drop (Waterproof), Cotton (Soft). 
      High-end branding style.`
    },
    // --- A+图 (3:2) ---
    {
      id: 'APlus_01_Banner',
      type: '品牌Banner',
      aspectRatio: '3:2',
      prompt: `Amazon A+ Banner. Wide shot of modern living room. 
      White cat wearing the ${va.material} cone walking on rug. 
      Cone flares outwards correctly. 
      Text space on left. Warm, inviting atmosphere.`
    },
    {
      id: 'APlus_02_Comparison',
      type: '竞品对比',
      aspectRatio: '3:2',
      prompt: `Comparison image. Split screen. 
      Left (Touchdog): Colorful, happy cat wearing soft blue cone. Green Checkmark. 
      Right (Others): Black and white, sad cat wearing hard plastic transparent cone. Red Cross. 
      Text: "Soft & Comfy" vs "Hard & Heavy".`
    },
    {
      id: 'APlus_03_Detail',
      type: '材质细节',
      aspectRatio: '3:2',
      prompt: `Extreme close-up macro shot of the product material. 
      Show the waterproof PU texture with water droplets beading off. 
      Show the neat stitching and soft edge binding. 
      High tech, quality feel.`
    },
    {
      id: 'APlus_04_Waterproof',
      type: '防水测试',
      aspectRatio: '3:2',
      prompt: `Action shot. Water being poured onto the blue cone collar. 
      Water repelling instantly, rolling off the surface. 
      Demonstrates "Easy to Wipe" feature. 
      Bright lighting.`
    },
    {
      id: 'APlus_05_Storage',
      type: '收纳展示',
      aspectRatio: '3:2',
      prompt: `Lifestyle shot. The cone collar folded/rolled up compactly. 
      Placed inside a small travel bag or drawer. 
      Shows "Easy Storage" feature. 
      Clean composition.`
    },
    {
      id: 'APlus_06_Guide',
      type: '选购指南',
      aspectRatio: '3:2',
      prompt: `Sizing guide banner. 
      Illustration of how to measure cat's neck circumference. 
      Touchdog branding elements. 
      Clean, instructional design.`
    }
  ];

  // 4. 执行生成
  console.log('\n🚀 开始全量生成 (12张)...');
  for (const item of prompts) {
    await generateSingleImage(item, refImageUrls);
  }
  
  console.log('\n✅ POC 完整验证结束。');
}

main().catch(console.error);