#!/usr/bin/env python3
"""
腾讯混元3D 图生3D 完整自动化脚本
使用 cloakbrowser 在浏览器内完成所有操作（包括签名生成）
"""

import json
import os
import time
from datetime import datetime
from cloakbrowser import launch_persistent_context

from ..config import get_profile_dir

PROFILE_DIR = str(get_profile_dir())


def generate_3d(image_path, wait_for_complete=False, timeout=300):
    """
    上传图片并触发生成3D模型

    Args:
        image_path: 图片路径
        wait_for_complete: 是否等待生成完成
        timeout: 最长等待时间（秒）

    Returns:
        dict: 包含 creationsId 和状态信息
    """
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"图片不存在: {image_path}")

    if not os.path.exists(PROFILE_DIR):
        raise FileNotFoundError(f"未找到登录状态目录: {PROFILE_DIR}")

    context = launch_persistent_context(PROFILE_DIR, headless=True)
    page = context.new_page()

    # 用于存储结果
    result = {"creationsId": None, "status": None, "modelUrl": None}

    # 监听响应
    def handle_response(response):
        url = response.url
        if "creations/generations" in url and response.status == 200:
            try:
                body = response.json()
                if "creationsId" in body:
                    result["creationsId"] = body["creationsId"]
                    print(f"✅ 生成已触发，creationsId: {body['creationsId']}")
            except:
                pass
        elif "creations/detail" in url and response.status == 200:
            try:
                body = response.json()
                result["status"] = body.get("status")
                if "result" in body and isinstance(body["result"], list) and len(body["result"]) > 0:
                    model_data = body["result"][0]
                    if "modelUrl" in model_data:
                        result["modelUrl"] = model_data["modelUrl"]
            except:
                pass

    page.on("response", handle_response)

    try:
        print("[1/6] 打开首页...")
        page.goto("https://3d.hunyuan.tencent.com/")
        page.wait_for_timeout(3000)

        print("[2/6] 点击 AI创作...")
        page.locator("button").filter(has_text="AI创作").first.click()
        page.wait_for_timeout(2000)

        print("[3/6] 点击 图生3D...")
        page.locator("text=图生3D").first.click()
        page.wait_for_timeout(2000)

        print("[4/6] 点击 单张图片...")
        page.locator("text=单张图片").first.click()
        page.wait_for_timeout(2000)

        print("[5/6] 上传图片...")
        page.locator('input[type=file]').first.set_input_files(image_path)
        page.wait_for_timeout(3000)

        print("[6/6] 点击 立即生成...")
        page.locator("text=立即生成").first.click()

        # 等待生成触发
        page.wait_for_timeout(5000)

        if not result["creationsId"]:
            print("⚠️ 未获取到 creationsId，可能生成失败")
            return result

        if wait_for_complete:
            print(f"\n⏳ 等待生成完成（最长 {timeout} 秒）...")
            start = time.time()
            last_status = None
            while time.time() - start < timeout:
                # 轮询状态
                status_result = page.evaluate(f'''
                async () => {{
                  const resp = await fetch('https://3d.hunyuan.tencent.com/api/3d/creations/detail?creationsId={result["creationsId"]}', {{
                    headers: {{'X-Source': 'web', 'Referer': 'https://3d.hunyuan.tencent.com/'}}
                  }});
                  return await resp.json();
                }}
                ''')

                status = status_result.get("status", "unknown")
                progress = status_result.get("progress", 0)

                if status != last_status:
                    print(f"  状态: {status} (进度: {progress}%)")
                    last_status = status

                if status == "success":
                    # 提取模型URL
                    if "result" in status_result and len(status_result["result"]) > 0:
                        model_data = status_result["result"][0]
                        result["modelUrl"] = model_data.get("modelUrl")
                        result["previewUrl"] = model_data.get("previewUrl")
                    print(f"\n✅ 生成完成！")
                    break
                elif status == "fail":
                    print(f"\n❌ 生成失败")
                    break

                time.sleep(3)
            else:
                print(f"\n⏰ 超时，当前状态: {last_status}")

        return result

    finally:
        context.close()


def get_quota():
    """获取当前配额信息"""
    context = launch_persistent_context(PROFILE_DIR, headless=True)
    page = context.new_page()

    try:
        page.goto("https://3d.hunyuan.tencent.com/")
        page.wait_for_timeout(3000)

        result = page.evaluate('''
        async () => {
          const resp = await fetch('https://3d.hunyuan.tencent.com/api/3d/quotainfo', {
            method: 'POST',
            headers: {'Content-Type': 'application/json', 'X-Source': 'web', 'Referer': 'https://3d.hunyuan.tencent.com/'},
            body: '{"sceneType":"3dCreations"}'
          });
          return await resp.json();
        }
        ''')
        return result
    finally:
        context.close()


def get_creations_list():
    """获取作品列表"""
    context = launch_persistent_context(PROFILE_DIR, headless=True)
    page = context.new_page()

    try:
        page.goto("https://3d.hunyuan.tencent.com/")
        page.wait_for_timeout(3000)

        result = page.evaluate('''
        async () => {
          const resp = await fetch('https://3d.hunyuan.tencent.com/api/3d/creations/list', {
            method: 'POST',
            headers: {'Content-Type': 'application/json', 'X-Source': 'web', 'Referer': 'https://3d.hunyuan.tencent.com/'},
            body: '{"page":1,"pageSize":20}'
          });
          return await resp.json();
        }
        ''')
        return result
    finally:
        context.close()


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("用法:")
        print(f"  python {sys.argv[0]} quota          # 查询配额")
        print(f"  python {sys.argv[0]} list           # 查询作品列表")
        print(f"  python {sys.argv[0]} generate <图片路径> [wait]  # 生成3D模型")
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "quota":
        quota = get_quota()
        print(json.dumps(quota, indent=2, ensure_ascii=False))

    elif cmd == "list":
        creations = get_creations_list()
        print(json.dumps(creations, indent=2, ensure_ascii=False))

    elif cmd == "generate":
        if len(sys.argv) < 3:
            print("请提供图片路径")
            sys.exit(1)

        image_path = sys.argv[2]
        wait = len(sys.argv) > 3 and sys.argv[3] == "wait"

        result = generate_3d(image_path, wait_for_complete=wait)
        print(json.dumps(result, indent=2, ensure_ascii=False))

    else:
        print(f"未知命令: {cmd}")