#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MULTI-REFERENCE-VIDEO-ADAPTER
多参考图视频适配器 — 支持苏白+牌匾+场景多参考输入。

功能:
1. 检查视频API是否支持多参考图输入
2. 如果支持: 封装多参考图接口，统一调用
3. 如果不支持: 明确报错，回退到"单参考图+后期合成"路线
4. 提供统一的调用接口给上游 Agent

用法:
  python multi-reference-video-adapter.py --prompt "苏白站在天道宗牌匾下" \\
    --references char-003-subai.png tdz-plaque.png env-baizonghui.png \\
    --output output.mp4

检查API能力:
  python multi-reference-video-adapter.py --check-api
"""

import os
import sys
import json
import argparse
import requests
from pathlib import Path
from datetime import datetime

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT / "engines"))

def load_api_config():
    """从苍耳本机仓库外密钥文件和环境变量加载配置"""
    config = {}
    env_files = [
        Path(os.environ["VIDEO_AI_SECRETS_FILE"]) if os.environ.get("VIDEO_AI_SECRETS_FILE") else None,
        Path("/root/guanghulab-local-secrets/cang-ying.env"),
        Path("~/guanghulab-local-secrets/cang-ying.env").expanduser(),
        Path("/Users/bingshuolingdianyuanhe/Documents/guanghulab-local-secrets/video-ai-system.env"),
    ]
    for env_file in env_files:
        if env_file is None:
            continue
        if not env_file.exists():
            continue
        with open(env_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#"):
                    key, _, val = line.partition("=")
                    config[key.strip()] = val.strip().strip('"').strip("'")
    for key in ["JIMENG_API_KEY", "JIMENG_BASE_URL", "JIMENG_MODEL"]:
        if os.environ.get(key):
            config[key] = os.environ[key]
    return config


class MultiReferenceVideoAdapter:
    """多参考图视频适配器"""

    def __init__(self, live_capability_check=False, allow_single_reference_fallback=False):
        self.config = load_api_config()
        self.api_key = self.config.get("JIMENG_API_KEY", "")
        self.base_url = self.config.get("JIMENG_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3")
        self.model = self.config.get("JIMENG_MODEL", "doubao-seedance-2-0-260128")
        self.live_capability_check = live_capability_check
        self.allow_single_reference_fallback = allow_single_reference_fallback

        # API 能力探测结果缓存
        self._api_capabilities = None

    def check_api_capabilities(self, live=None):
        """
        检查 API 是否支持多参考图
        返回: {
            "multi_reference_supported": bool,
            "max_references": int,
            "supported_types": list,  # ["image_url", "image_url_2", ...]
            "details": str
        }
        """
        if self._api_capabilities:
            return self._api_capabilities

        live = self.live_capability_check if live is None else live
        print("🔍 检查 Seedance API 多参考图支持...")

        # 根据 Volcengine 官方文档 (https://www.volcengine.com/docs/82379/1520757)
        # Seedance 2.0 API 的 content 数组支持多个 image_url 对象
        # 但需要实际测试确认

        # 理论上的 API 结构:
        # content: [
        #   { type: "text", text: "..." },
        #   { type: "image_url", image_url: { url: "data:image/png;base64,..." } },
        #   { type: "image_url", image_url: { url: "data:image/png;base64,..." } },  # 第二张参考图
        # ]

        if not live:
            test_result = {
                "multi_reference_supported": False,
                "max_references": 1,
                "supported_types": ["image_url"],
                "status": "unverified",
                "details": "未执行付费/联网实测；需要 --live-check 才会提交真实 API 探测任务",
                "requires_live_check": True
            }
        else:
            test_result = self._test_multi_reference()

        self._api_capabilities = test_result
        return test_result

    def _test_multi_reference(self):
        """
        实际测试 API 是否支持多参考图
        方法: 提交一个测试请求，包含2张参考图，观察响应
        """
        if not self.api_key:
            return {
                "multi_reference_supported": False,
                "max_references": 1,
                "supported_types": ["image_url"],
                "details": "未找到 JIMENG_API_KEY，无法执行 live check",
                "status": "missing_api_key"
            }

        # 构造一个最小测试请求
        test_prompt = "test multi-reference support"

        # 创建1x1像素的测试图片 (PNG)
        import base64
        from io import BytesIO
        try:
            from PIL import Image
            img = Image.new("RGB", (32, 32), color=(255, 0, 0))
            buf = BytesIO()
            img.save(buf, format="PNG")
            test_img_b64 = base64.b64encode(buf.getvalue()).decode()
        except ImportError:
            # 如果没有 PIL，用空base64
            test_img_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWg"

        # 构造 content 数组 (2张参考图)
        content = [
            {"type": "text", "text": test_prompt},
            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{test_img_b64}"}},
            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{test_img_b64}"}},
        ]

        payload = {
            "model": self.model,
            "content": content,
            "duration": 4,  # 最短时长，省钱
            "resolution": "480p"
        }

        # 发送请求
        try:
            print("  📤 发送测试请求 (2张参考图)...")
            url = f"{self.base_url}/contents/generations/tasks"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }

            response = requests.post(url, json=payload, headers=headers, timeout=30)

            if response.status_code == 200:
                # 成功！API 支持多参考图
                print("  ✅ API 支持多参考图！")
                return {
                    "multi_reference_supported": True,
                    "max_references": 2,  # 需要逐步测试确定上限
                    "supported_types": ["image_url"],
                    "details": "API 成功接受2张参考图"
                }
            elif response.status_code == 400:
                # 看错误信息
                error_data = response.json()
                error_msg = error_data.get("error", {}).get("message", "")
                print(f"  ❌ API 不支持多参考图: {error_msg}")
                return {
                    "multi_reference_supported": False,
                    "max_references": 1,
                    "supported_types": ["image_url"],  # 只支持单张
                    "details": error_msg,
                    "error_response": error_data
                }
            else:
                print(f"  ⚠️  未知响应: {response.status_code}")
                return {
                    "multi_reference_supported": False,
                    "max_references": 1,
                    "details": f"Unknown response: {response.status_code}"
                }

        except Exception as e:
            print(f"  ❌ 测试失败: {e}")
            return {
                "multi_reference_supported": False,
                "max_references": 1,
                "details": f"Test failed: {e}"
            }

    def generate_video(self, prompt, reference_images, output_path=None, duration=5, resolution="720p"):
        """
        生成视频 (多参考图)

        参数:
          prompt: str - 提示词
          reference_images: list[str] - 参考图路径列表
          output_path: str - 输出路径
          duration: int - 时长 (4-15)
          resolution: str - 分辨率 ("480p" | "720p")

        返回:
          {
            "success": bool,
            "task_id": str,
            "output_path": str,
            "method": str,  # "multi-reference" | "single-reference+composite"
            "warning": str
          }
        """
        print(f"\n🎬 生成视频 (多参考图)")
        print(f"   提示词: {prompt[:60]}...")
        print(f"   参考图数量: {len(reference_images)}")
        for i, img in enumerate(reference_images):
            print(f"      [{i+1}] {Path(img).name}")

        # 检查 API 能力
        capabilities = self.check_api_capabilities()

        if capabilities["multi_reference_supported"]:
            # API 支持多参考图 → 直接调用
            print(f"\n  ✅ API 支持多参考图，直接调用...")
            result = self._generate_multi_reference(prompt, reference_images, output_path, duration, resolution)
            result["method"] = "multi-reference"
            return result
        else:
            # API 不支持多参考图 → 明确报错 + 建议回退方案
            print(f"\n  ❌ API 不支持多参考图")
            print(f"  📋 错误详情: {capabilities['details']}")
            print(f"\n  💡 回退方案:")
            print(f"     1. 使用第一张参考图 (苏白) 生成视频")
            print(f"     2. 后期合成牌匾/场景 (平面追踪 + 贴图)")
            print(f"     3. 或使用可灵生成角色，Seedance 生成场景，后期合成")

            if not self.allow_single_reference_fallback:
                return {
                    "success": False,
                    "method": "not-run",
                    "error": "多参考图能力未确认或不支持，未自动回退单参考图以避免误导一致性测试",
                    "capabilities": capabilities,
                    "next_step": "使用 --live-check 验证多参考图，或显式传入 --allow-single-reference-fallback 做调试占位"
                }

            # 回退: 只用第一张参考图。必须显式开启，避免误导牌匾/场景一致性判断。
            warning = "API不支持多参考图，已回退到单参考图模式。牌匾/场景一致性需要后期合成。"
            print(f"\n  🔄 回退: 使用第一张参考图生成...")

            result = self._generate_single_reference(prompt, reference_images[0], output_path, duration, resolution)
            result["method"] = "single-reference+composite"
            result["warning"] = warning
            result["fallback_reason"] = capabilities["details"]

            return result

    def _generate_multi_reference(self, prompt, reference_images, output_path, duration, resolution):
        """调用多参考图 API"""
        # 构造 content 数组
        content = [{"type": "text", "text": prompt}]

        for img_path in reference_images:
            img_path = Path(img_path)
            if not img_path.exists():
                print(f"  ⚠️  参考图不存在: {img_path}")
                continue

            # 读取图片并转 base64
            import base64
            with open(img_path, "rb") as f:
                img_data = f.read()
            b64 = base64.b64encode(img_data).decode()
            mime = "image/png" if img_path.suffix.lower() == ".png" else "image/jpeg"

            content.append({
                "type": "image_url",
                "image_url": {"url": f"data:{mime};base64,{b64}"}
            })

        # 调用 API
        payload = {
            "model": self.model,
            "content": content,
            "duration": duration,
            "resolution": resolution
        }

        print(f"  📤 提交任务...")
        url = f"{self.base_url}/contents/generations/tasks"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        response = requests.post(url, json=payload, headers=headers, timeout=60)
        response.raise_for_status()
        data = response.json()

        task_id = data.get("id") or data.get("task_id")
        print(f"  ✅ 任务已提交: {task_id}")

        # 返回任务ID，等待轮询
        return {
            "success": True,
            "task_id": task_id,
            "output_path": output_path,
            "api_response": data
        }

    def _generate_single_reference(self, prompt, reference_image, output_path, duration, resolution):
        """回退: 单参考图生成"""
        # 调用现有的 video-api-adapter (Node.js)
        # 这里用 subprocess 调用
        import subprocess

        print(f"  📤 调用单参考图 API...")

        # 构造调用参数
        node_script = PROJECT_ROOT / "engines" / "video-api-adapter.js"
        cmd = [
            "node", str(node_script),
            "--prompt", prompt,
            "--reference-image", str(reference_image),
            "--duration", str(duration),
            "--resolution", resolution
        ]

        # 执行
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            print(f"  ❌ 调用失败: {result.stderr}")
            return {"success": False, "error": result.stderr}

        print(f"  ✅ 任务已提交")
        return {"success": True, "method": "single-reference", "stdout": result.stdout}

    def batch_generate(self, shots_config):
        """
        批量生成 (从配置文件)

        shots_config 格式:
          [
            {
              "shot_id": "E1-SHOT01",
              "prompt": "...",
              "references": ["char.png", "prop.png", "env.png"],
              "output": "output/E1-SHOT01.mp4"
            },
            ...
          ]
        """
        results = []
        for shot in shots_config:
            result = self.generate_video(
                prompt=shot["prompt"],
                reference_images=shot["references"],
                output_path=shot["output"]
            )
            results.append(result)
        return results


def main():
    parser = argparse.ArgumentParser(description="MULTI-REFERENCE-VIDEO-ADAPTER")
    parser.add_argument("--check-api", action="store_true", help="检查API多参考图支持")
    parser.add_argument("--live-check", action="store_true",
                        help="执行真实 API 探测请求，可能产生费用")
    parser.add_argument("--allow-single-reference-fallback", action="store_true",
                        help="多参考图不可用时允许单参考图调试占位")
    parser.add_argument("--prompt", type=str, help="提示词")
    parser.add_argument("--references", type=str, nargs="+", help="参考图路径列表")
    parser.add_argument("--output", type=str, help="输出路径")
    parser.add_argument("--duration", type=int, default=5, help="时长 (4-15)")
    parser.add_argument("--resolution", type=str, default="720p", choices=["480p", "720p"], help="分辨率")

    args = parser.parse_args()

    adapter = MultiReferenceVideoAdapter(
        live_capability_check=args.live_check,
        allow_single_reference_fallback=args.allow_single_reference_fallback
    )

    if args.check_api:
        capabilities = adapter.check_api_capabilities(live=args.live_check)
        print(f"\n📊 API 能力报告:")
        print(f"   多参考图支持: {capabilities['multi_reference_supported']}")
        print(f"   最大参考图数: {capabilities['max_references']}")
        print(f"   详情: {capabilities['details']}")
        return

    if not args.prompt or not args.references:
        parser.print_help()
        return

    result = adapter.generate_video(
        prompt=args.prompt,
        reference_images=args.references,
        output_path=args.output,
        duration=args.duration,
        resolution=args.resolution
    )

    print(f"\n📋 生成结果:")
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()