#!/usr/bin/env node

/**
 * Kling Image Generator - 可灵多图参考生图 (风格参考)
 *
 * 使用 /v1/images/multi-image2image 端点，支持 style_image 风格参考
 * 配置来源: config.json 的 kelingApiKey + kelingSecretAccessKey + kelingApiBaseUrl
 *
 * 用法:
 *   node kling-image-generator.js "a cute cat" -r http://img.com/style.jpg -o ./output
 */

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

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

const configPath = path.join(__dirname, '..', '..', 'config.json')
const cfg = fs.existsSync(configPath) ? JSON.parse(fs.readFileSync(configPath, 'utf-8')) : {}

const Config = {
  apiBase: (cfg.kelingApiBaseUrl || 'https://api-beijing.klingai.com').replace(/\/+$/, ''),
  ak: (cfg.kelingApiKey || '').trim(),
  sk: (cfg.kelingSecretAccessKey || '').trim(),
  model: cfg.kelingImageModel || 'kling-v2-1',
  pollInterval: 5000,
  maxPollTime: 300000,
}

// ============================================================================
// JWT 鉴权
// ============================================================================

function base64url(buf) {
  return Buffer.from(buf).toString('base64')
    .replace(/=/g, '').replace(/\+/g, '-').replace(/\//g, '_')
}

function getToken() {
  if (!Config.ak || !Config.sk) throw new Error('未配置可灵 API 凭证（config.json 需要 kelingApiKey + kelingSecretAccessKey）')
  const header = base64url(JSON.stringify({ alg: 'HS256', typ: 'JWT' }))
  const now = Math.floor(Date.now() / 1000)
  const payload = base64url(JSON.stringify({ iss: Config.ak, exp: now + 1800, nbf: now - 5 }))
  const sig = base64url(crypto.createHmac('sha256', Config.sk).update(`${header}.${payload}`).digest())
  return `${header}.${payload}.${sig}`
}

// ============================================================================
// 图片下载
// ============================================================================

async function download(url, outputPath) {
  const protocol = url.startsWith('https') ? https : http
  return new Promise((resolve, reject) => {
    const file = fs.createWriteStream(outputPath)
    protocol.get(url, (res) => {
      if (res.statusCode >= 300 && res.statusCode < 400 && res.headers.location) {
        file.close()
        fs.unlinkSync(outputPath)
        return download(res.headers.location, outputPath).then(resolve).catch(reject)
      }
      res.pipe(file)
      file.on('finish', () => { file.close(); resolve(outputPath) })
    }).on('error', (err) => {
      file.close()
      if (fs.existsSync(outputPath)) fs.unlinkSync(outputPath)
      reject(err)
    })
  })
}

// ============================================================================
// 可灵多图参考生图 API (/v1/images/multi-image2image)
// ============================================================================

const KlingImageApi = {
  async submit(prompt, options = {}) {
    const { styleImageUrl = '', aspectRatio = '' } = options
    const token = getToken()

    const body = {
      model_name: Config.model,
      prompt,
      n: 1,
    }
    if (styleImageUrl) {
      body.subject_image_list = [{ subject_image: styleImageUrl }]
      body.style_image = styleImageUrl
    }
    if (aspectRatio) body.aspect_ratio = aspectRatio

    console.log(`\n📡 提交可灵图片任务`)
    console.log(`   模型: ${Config.model}`)
    console.log(`   提示词: ${prompt.substring(0, 80)}...`)
    if (styleImageUrl) console.log(`   风格图: ${styleImageUrl.substring(0, 60)}...`)

    const res = await fetch(`${Config.apiBase}/v1/images/multi-image2image`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${token}` },
      body: JSON.stringify(body),
    })

    const json = await res.json()
    if (json.code !== undefined && json.code !== 0 && json.code !== 200) {
      throw new Error(`可灵图片 API 错误 (code=${json.code}): ${json.message || JSON.stringify(json)}`)
    }

    const taskId = (json.data || {}).task_id
    if (!taskId) throw new Error(`可灵图片提交失败: ${JSON.stringify(json)}`)

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

  async poll(taskId) {
    const startTime = Date.now()
    console.log(`\n⏳ 等待可灵生成...`)

    while (Date.now() - startTime < Config.maxPollTime) {
      const token = getToken()
      const res = await fetch(`${Config.apiBase}/v1/images/multi-image2image/${taskId}`, {
        headers: { 'Authorization': `Bearer ${token}` },
      })

      const json = await res.json()
      const data = json.data || {}
      const status = data.task_status

      if (status === 'succeed') {
        const images = data.task_result?.images || []
        if (images.length > 0) {
          console.log(`\n✅ 可灵图片生成完成!`)
          return { success: true, imageUrl: images[0].url }
        }
        throw new Error('可灵返回成功但无图片')
      }

      if (status === 'failed') {
        throw new Error(`可灵图片生成失败: ${data.task_status_msg || '未知原因'}`)
      }

      process.stdout.write(`   状态: ${status || 'processing'}\r`)
      await new Promise(r => setTimeout(r, Config.pollInterval))
    }

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

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

async function generate(prompt, options = {}) {
  const { outputDir = './output', styleImageUrl = '', aspectRatio = '' } = options

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

  const taskId = await KlingImageApi.submit(prompt, { styleImageUrl, aspectRatio })
  const result = await KlingImageApi.poll(taskId)

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

  return { success: true, taskId, imageUrl: result.imageUrl, savedFiles: [filePath] }
}

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

async function main() {
  const args = process.argv.slice(2)
  if (args.includes('-h') || args.includes('--help') || args.length === 0) {
    console.log(`
用法: node kling-image-generator.js <prompt> [options]

选项:
  -o, --output <dir>  输出目录 (默认: ./output)
  -r, --ref <url>     风格参考图 URL
  -a, --ar <ratio>    画幅比例 (如 9:16, 16:9)
  -h, --help          帮助
`)
    return
  }

  const options = { outputDir: './output', styleImageUrl: '', aspectRatio: '' }
  const params = []

  let i = 0
  while (i < args.length) {
    const arg = args[i]
    if (arg === '-o' || arg === '--output') { options.outputDir = args[++i] }
    else if (arg === '-r' || arg === '--ref') { options.styleImageUrl = args[++i] }
    else if (arg === '-a' || arg === '--ar') { options.aspectRatio = args[++i] }
    else { params.push(arg) }
    i++
  }

  const prompt = params.join(' ')
  if (!prompt) { console.error('请提供提示词'); process.exit(1) }
  await generate(prompt, options)
}

module.exports = { generate, KlingImageApi }

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