#!/usr/bin/env node

/**
 * MJ Image Generator - Midjourney 图片生成工具
 *
 * 功能：
 * - 提交 imagine 任务
 * - 轮询直到完成
 * - 下载 4 合 1 结果图
 * - 自动拆分为 4 张独立图片
 * - 支持参考图（图生图）
 *
 * 用法：
 *   node mj-image-generator.js "a cute cat" -o ./output
 *   node mj-image-generator.js "cyberpunk city" -o ./output -r http://example.com/ref.jpg
 *   node mj-image-generator.js batch ./prompts.txt -o ./output
 */

const fs = require('fs')
const path = require('path')
const https = require('https')
const http = require('http')
const sharp = require('sharp')

// ============================================================================
// 配置
// ============================================================================

function loadConfig() {
  const configPath = path.join(__dirname, '..', '..', 'config.json')
  if (fs.existsSync(configPath)) {
    return JSON.parse(fs.readFileSync(configPath, 'utf-8'))
  }
  return {}
}

const cfg = loadConfig()

const Config = {
  baseUrl: cfg.mjApiBaseUrl,
  apiKey: cfg.mjApiKey || '',
  pollInterval: 5000,     // 轮询间隔 5 秒
  maxPollTime: 300000,    // 最大等待 5 分钟
}

// ============================================================================
// API 调用
// ============================================================================

const MJApi = {
  /**
   * 提交 imagine 任务
   */
  async submit(prompt, options = {}) {
    const { referenceImages = [], botType = 'mj', aspectRatio = '', mjParams = '', styleWeight = 100 } = options

    let finalPrompt = prompt
    if (referenceImages.length > 0) {
      const srefSection = `--sref ${referenceImages.join(' ')} --sw ${styleWeight}`
      finalPrompt = `${prompt} ${srefSection}`.trim()
    }
    if (aspectRatio) {
      finalPrompt = `${finalPrompt} --ar ${aspectRatio}`
    }
    if (mjParams) {
      finalPrompt = `${finalPrompt} ${mjParams}`
    }

    const body = {
      prompt: finalPrompt,
      base64Array: [],
      botType,
    }

    console.log(`\n📡 提交 MJ 任务`)
    console.log(`   提示词: ${finalPrompt.substring(0, 80)}...`)
    console.log(`   参考图: ${referenceImages.length} 张`)

    const res = await fetch(`${Config.baseUrl}/mj/submit/imagine`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${Config.apiKey}`,
      },
      body: JSON.stringify(body),
    })

    const result = await res.json()

    if (result.code !== 1 && result.code !== '1') {
      throw new Error(`MJ 提交失败: code=${result.code}, ${result.description || JSON.stringify(result)}`)
    }

    const taskId = result.result
    console.log(`   任务 ID: ${taskId}`)
    return taskId
  },

  /**
   * 轮询任务状态
   */
  async poll(taskId) {
    const startTime = Date.now()
    let lastProgress = ''

    console.log(`\n⏳ 等待 MJ 生成...`)

    while (Date.now() - startTime < Config.maxPollTime) {
      const res = await fetch(`${Config.baseUrl}/mj/task/${taskId}/fetch`, {
        headers: { 'Authorization': `Bearer ${Config.apiKey}` },
      })

      const task = await res.json()
      const status = task.status

      if (status === 'SUCCESS') {
        console.log(`\n✅ 生成完成!`)
        console.log(`   图片 URL: ${task.imageUrl}`)
        return {
          success: true,
          imageUrl: task.imageUrl,
          prompt: task.prompt || task.promptEn,
        }
      }

      if (status === 'FAILURE') {
        const errMsg = task.failReason || '未知原因'
        throw new Error(`MJ 生成失败: ${errMsg}`)
      }

      // 显示进度
      const progress = task.progress || ''
      if (progress !== lastProgress) {
        lastProgress = progress
        process.stdout.write(`   进度: ${progress}%   状态: ${status}\r`)
      }

      await new Promise(r => setTimeout(r, Config.pollInterval))
    }

    throw new Error(`MJ 生成超时 (${Config.maxPollTime / 1000}s)`)
  },
}

// ============================================================================
// 图片处理
// ============================================================================

const ImageUtils = {
  /**
   * 下载图片到本地
   */
  async download(url, outputPath) {
    const protocol = url.startsWith('https') ? https : http

    return new Promise((resolve, reject) => {
      const file = fs.createWriteStream(outputPath)
      protocol.get(url, (response) => {
        // 处理重定向
        if (response.statusCode >= 300 && response.statusCode < 400 && response.headers.location) {
          file.close()
          fs.unlinkSync(outputPath)
          return ImageUtils.download(response.headers.location, outputPath).then(resolve).catch(reject)
        }
        response.pipe(file)
        file.on('finish', () => {
          file.close()
          resolve(outputPath)
        })
      }).on('error', (err) => {
        file.close()
        if (fs.existsSync(outputPath)) fs.unlinkSync(outputPath)
        reject(err)
      })
    })
  },

  /**
   * 将 4 合 1 图片拆分为 4 张独立图片
   */
  async split4(gridImagePath, outputDir, prefix = 'mj') {
    const image = sharp(gridImagePath)
    const metadata = await image.metadata()
    const { width, height } = metadata

    // MJ 4 合 1 是 2x2 网格，每格约一半
    const halfW = Math.floor(width / 2)
    const halfH = Math.floor(height / 2)

    const positions = [
      { name: `${prefix}_1`, x: 0, y: 0 },           // 左上
      { name: `${prefix}_2`, x: halfW, y: 0 },        // 右上
      { name: `${prefix}_3`, x: 0, y: halfH },        // 左下
      { name: `${prefix}_4`, x: halfW, y: halfH },    // 右下
    ]

    const files = []
    for (const pos of positions) {
      const outputPath = path.join(outputDir, `${pos.name}.png`)
      await sharp(gridImagePath)
        .extract({ left: pos.x, top: pos.y, width: halfW, height: halfH })
        .toFile(outputPath)
      files.push(outputPath)
    }

    return files
  },
}

// ============================================================================
// 核心流程
// ============================================================================

async function generate(prompt, options = {}) {
  const {
    outputDir = './output',
    referenceImages = [],
    aspectRatio = '',
    mjParams = '',
    split = true,
    keepGrid = false,
  } = options

  if (!Config.apiKey) {
    throw new Error('未配置 mjApiKey，请在 config.json 中添加')
  }

  fs.mkdirSync(outputDir, { recursive: true })

  // 1. 提交任务
  const taskId = await MJApi.submit(prompt, { referenceImages, aspectRatio, mjParams, styleWeight: options.styleWeight })

  // 2. 轮询等待
  const result = await MJApi.poll(taskId)

  // 3. 下载
  const timestamp = new Date().toISOString().replace(/[:.]/g, '-')
  const gridFile = path.join(outputDir, `${timestamp}_grid.png`)
  await ImageUtils.download(result.imageUrl, gridFile)
  console.log(`   下载完成: ${gridFile}`)

  // 4. 拆分
  const allFiles = [gridFile]
  if (split) {
    const prefix = timestamp
    const splitFiles = await ImageUtils.split4(gridFile, outputDir, prefix)
    allFiles.push(...splitFiles)
    console.log(`   拆分完成: ${splitFiles.length} 张图片`)

    if (!keepGrid) {
      fs.unlinkSync(gridFile)
      allFiles.shift()
      console.log(`   已删除网格图`)
    }
  }

  return {
    success: true,
    taskId,
    imageUrl: result.imageUrl,
    files: allFiles,
  }
}

/**
 * 批量生成（并发提交 + 并行轮询）
 */
async function batchGenerate(prompts, options = {}) {
  const concurrency = options.concurrency || prompts.length
  const results = new Array(prompts.length).fill(null)

  // 分批提交，每批 concurrency 个并行
  for (let batchStart = 0; batchStart < prompts.length; batchStart += concurrency) {
    const batchEnd = Math.min(batchStart + concurrency, prompts.length)
    const batchIndices = []
    for (let i = batchStart; i < batchEnd; i++) batchIndices.push(i)

    // 并行提交所有任务
    console.log(`\n📡 批量提交 [${batchStart + 1}-${batchEnd}/${prompts.length}]...`)
    const taskIds = await Promise.all(batchIndices.map(async (i) => {
      try {
        const taskId = await MJApi.submit(prompts[i], {
          referenceImages: options.referenceImages,
          aspectRatio: options.aspectRatio,
          mjParams: options.mjParams,
          styleWeight: options.styleWeight,
        })
        return { i, taskId }
      } catch (err) {
        console.error(`   [${i + 1}] ❌ 提交失败: ${err.message}`)
        results[i] = { success: false, prompt: prompts[i], error: err.message }
        return { i, taskId: null }
      }
    }))

    // 过滤出成功的任务，并行轮询
    const activeTasks = taskIds.filter(t => t.taskId)
    console.log(`   已提交 ${activeTasks.length} 个任务，并行等待生成...\n`)

    const pollResults = await Promise.all(activeTasks.map(async ({ i, taskId }) => {
      try {
        const result = await MJApi.poll(taskId)
        // 下载 + 拆分
        const timestamp = new Date().toISOString().replace(/[:.]/g, '-')
        const gridFile = path.join(options.outputDir || './output', `${timestamp}_grid.png`)
        fs.mkdirSync(path.dirname(gridFile), { recursive: true })
        await ImageUtils.download(result.imageUrl, gridFile)

        const allFiles = [gridFile]
        if (options.split !== false) {
          const prefix = timestamp
          const splitFiles = await ImageUtils.split4(gridFile, options.outputDir || './output', prefix)
          allFiles.push(...splitFiles)
          if (!options.keepGrid) {
            fs.unlinkSync(gridFile)
            allFiles.shift()
          }
        }

        console.log(`   [${i + 1}/${prompts.length}] ✅ 完成`)
        return { i, success: true, prompt: prompts[i], taskId, imageUrl: result.imageUrl, files: allFiles }
      } catch (err) {
        console.error(`   [${i + 1}/${prompts.length}] ❌ 失败: ${err.message}`)
        return { i, success: false, prompt: prompts[i], error: err.message }
      }
    }))

    for (const r of pollResults) results[r.i] = r
  }

  const ok = results.filter(r => r && r.success).length
  console.log(`\n✨ 批量完成: ${ok}/${prompts.length} 成功`)

  return results
}

// ============================================================================
// CLI
// ============================================================================

function showHelp() {
  console.log(`
🎨 MJ Image Generator - Midjourney 图片生成工具

用法:
  node mj-image-generator.js <prompt> [options]
  node mj-image-generator.js batch <file> [options]

选项:
  -o, --output <dir>       输出目录 (默认: ./output)
  -r, --ref <urls>         参考图 URL，逗号分隔
  -a, --ar <ratio>         宽高比 (1:1, 16:9, 9:16, 3:4, 4:3 等)
  -c, --concurrency <n>    并发数 (默认: 全部并行)
  --no-split               不拆分 4 合 1
  --keep-grid              保留原始网格图
  -h, --help               帮助

示例:
  node mj-image-generator.js "a cute cat" -o ./cats
  node mj-image-generator.js "cyberpunk city" -a 16:9
  node mj-image-generator.js "portrait" -r http://img.com/ref.jpg -a 9:16
  node mj-image-generator.js batch ./prompts.txt -o ./batch
`)
}

async function main() {
  const args = process.argv.slice(2)

  if (args.includes('-h') || args.includes('--help') || args.length === 0) {
    showHelp()
    return
  }

  let command = 'generate'
  let params = []
  const options = { outputDir: './output', split: true, keepGrid: false, referenceImages: [], aspectRatio: '', concurrency: 0, mjParams: '', styleWeight: 100 }

  let i = 0
  if (args[0] === 'batch') {
    command = 'batch'
    i = 1
  }

  while (i < args.length) {
    const arg = args[i]
    if (arg === '-o' || arg === '--output') {
      options.outputDir = args[++i]
    } else if (arg === '-a' || arg === '--ar') {
      options.aspectRatio = args[++i]
    } else if (arg === '-r' || arg === '--ref') {
      options.referenceImages = args[++i].split(',').map(s => s.trim()).filter(Boolean)
    } else if (arg === '--no-split') {
      options.split = false
    } else if (arg === '--keep-grid') {
      options.keepGrid = true
    } else if (arg === '-c' || arg === '--concurrency') {
      options.concurrency = parseInt(args[++i], 10) || 0
    } else if (arg === '--mj-params') {
      options.mjParams = args[++i]
    } else if (arg === '--sw') {
      options.styleWeight = parseInt(args[++i], 10) || 100
    } else {
      params.push(arg)
    }
    i++
  }

  if (command === 'batch') {
    const filePath = params[0]
    if (!filePath || !fs.existsSync(filePath)) {
      console.error('请提供提示词文件路径')
      process.exit(1)
    }
    const prompts = fs.readFileSync(filePath, 'utf-8')
      .split('\n').filter(l => l.trim()).map(l => l.trim())
    await batchGenerate(prompts, options)
  } else {
    const prompt = params.join(' ')
    if (!prompt) {
      console.error('请提供提示词')
      process.exit(1)
    }
    await generate(prompt, options)
  }
}

// ============================================================================
// 导出
// ============================================================================

module.exports = { generate, batchGenerate, MJApi, ImageUtils }

if (require.main === module) {
  main().catch(err => {
    console.error(`\n❌ 错误: ${err.message}`)
    process.exit(1)
  })
}