#!/usr/bin/env python3
"""
腾讯混元3D 完整API客户端
支持所有生成模式：图生3D、文生3D、草图生3D、动画生成、纹理生成、智能拓扑
"""

import os
import requests
import json
import time
from typing import Optional, Dict, Any, List
from urllib.parse import urlparse, unquote
from .sign import sign

BASE_URL = "https://3d.hunyuan.tencent.com"
API_BASE = f"{BASE_URL}/api/3d"


class Hunyuan3DAPI:
    """腾讯混元3D 完整API客户端"""

    # 生成模式常量
    MODEL_TYPE_IMAGE = "image2ModelV3.1"
    MODEL_TYPE_TEXT = "text2ModelV3.1"
    MODEL_TYPE_SKETCH = "sketch2ModelV3.1"
    MODEL_TYPE_ANIMATION = "animation3dV2"
    MODEL_TYPE_TEXTURE = "textureTo3DV2"
    MODEL_TYPE_LOWPOLY = "lowpolyV2"

    # 动作类型
    MOTION_CAPOEIRA = 9
    MOTION_FALLING = 10
    MOTION_JUMPING = 11
    MOTION_KICKING = 12
    MOTION_SWORD = 13
    MOTION_RUNNING = 15
    MOTION_DANCING = 16

    # 纹理风格
    STYLE_DEFAULT = ""
    STYLE_SCULPTURE = "sculpture"
    STYLE_QINGHUA = "qinghuaci"
    STYLE_CHINA = "china_style"
    STYLE_CARTOON = "cartoon"
    STYLE_CYBERPUNK = "cyberpunk"

    # 拓扑面数
    TOPO_LOW = 5000
    TOPO_MEDIUM = 18000
    TOPO_HIGH = 30000

    # 模型下载格式键名
    # 原生格式：直接来自 API urlResult
    # 转换格式：通过 /creations/resourceConvert 从 OBJ(zip) 转换得到
    MODEL_FORMAT_KEYS = [
        "glb",          # 带PBR贴图的GLB二进制模型
        "obj",          # OBJ模型（通常打包在zip中）
        "mtl",          # MTL材质文件
        "obj_url",      # 独立OBJ文件URL
        "geometryGlb",  # 纯几何GLB（无材质）
        "textureGlb",   # GLB纹理
        "textureObj",   # OBJ纹理包
        "image_url",    # 预览图/缩略图
        "pbrImage",           # PBR综合贴图
        "pbrMetallicImage",   # PBR金属度贴图
        "pbrRoughnessImage",  # PBR粗糙度贴图
        "pbrNormalImage",     # PBR法线贴图
        "invisible_wall",     # 不可见碰撞墙
        "air_wall",           # 空气墙（碰撞体）
        "fbx",          # FBX 格式（转换）
        "stl",          # STL 格式（转换）
        "usdz",         # USDZ 格式（转换，用于 iOS AR）
        "mp4",          # MP4 视频格式（转换）
        "gif",          # GIF 动图格式（转换）
    ]

    def __init__(self, cookies: Optional[str] = None):
        self.session = requests.Session()
        self.session.headers.update({
            "Content-Type": "application/json",
            "x-source": "web",
            "x-product": "hunyuan3d",
            "Origin": BASE_URL,
            "Referer": f"{BASE_URL}/",
            "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
        })
        try:
            cookie_value = cookies if cookies is not None else load_cookies_from_file()
        except FileNotFoundError:
            cookie_value = None
        if cookie_value:
            self.set_cookies(cookie_value)

    def set_cookies(self, cookies: str):
        self.session.headers["Cookie"] = cookies

    def _make_request(self, method: str, endpoint: str, params: Optional[Dict] = None,
                      data: Optional[Dict] = None) -> Dict[str, Any]:
        url = f"{API_BASE}{endpoint}"
        sign_params = dict(params) if params else {}
        signed_params = sign(sign_params)

        timestamp = signed_params.pop("timestamp")
        nonce = signed_params.pop("nonce")
        sign_value = signed_params.pop("sign")

        query_params = {}
        if params:
            query_params.update(params)
        query_params.update({
            "timestamp": timestamp,
            "nonce": nonce,
            "sign": sign_value
        })

        if method.upper() == "GET":
            response = self.session.get(url, params=query_params, timeout=30)
        else:
            response = self.session.post(url, params=query_params, json=data, timeout=30)

        response.raise_for_status()
        return response.json()

    # ========== 用户接口 ==========

    def get_user_info(self) -> Dict[str, Any]:
        """获取用户信息"""
        return self._make_request("GET", "/getuserinfo")

    def get_quota_info(self, scene_type: str = "3dCreations") -> Dict[str, Any]:
        """获取配额信息"""
        return self._make_request("POST", "/quotainfo", data={"sceneType": scene_type})

    # ========== 作品管理 ==========

    def get_creation_list(self, page: int = 1, page_size: int = 20) -> Dict[str, Any]:
        """获取作品列表"""
        return self._make_request("POST", "/creations/list", data={
            "page": page,
            "pageSize": page_size
        })

    def get_creation_count(self, status_list: Optional[List[str]] = None) -> Dict[str, Any]:
        """获取作品数量"""
        data = {"statusList": status_list or ["wait", "processing", "success", "fail"]}
        return self._make_request("POST", "/creations/count", data=data)

    def get_generation_status(self, creation_id: str) -> Dict[str, Any]:
        """获取生成状态"""
        return self._make_request("GET", "/creations/detail", params={
            "creationsId": creation_id
        })

    def resource_convert(self, source_url: str, target_formats: List[str]) -> Dict[str, str]:
        """
        调用资源转换接口将 OBJ(zip) 转换为其他格式

        Args:
            source_url: OBJ zip 文件的 URL（来自 urlResult['obj']）
            target_formats: 目标格式列表，如 ["usdz", "fbx"]

        Returns:
            格式键名 -> 转换后下载URL 的字典
        """
        data = {
            "sourceResource": [
                {"format": "zip", "url": source_url}
            ],
            "targetFormatList": target_formats
        }
        resp = self._make_request("POST", "/creations/resourceConvert", data=data)
        result = {}
        for item in resp.get("convertResult", []):
            fmt = item.get("format")
            url = item.get("url")
            if fmt and url:
                result[fmt] = url
        return result

    def get_model_urls(self, creation_id: str, include_converted: bool = False) -> Dict[str, Optional[str]]:
        """
        获取指定创作所有可用格式的下载URL

        Args:
            creation_id: 创作ID
            include_converted: 是否包含转换格式（fbx/stl/usdz/mp4/gif），
                               启用时会调用 resourceConvert 接口

        Returns:
            格式键名 -> 下载URL 的字典，空值已被过滤
        """
        status = self.get_generation_status(creation_id)

        # 兼容两种响应包装格式
        if "code" in status:
            data = status.get("data", {})
        else:
            data = status

        result_list = data.get("result", [])
        if not result_list:
            return {}

        url_result = result_list[0].get("urlResult", {})
        urls = {
            key: val
            for key in self.MODEL_FORMAT_KEYS
            if (val := url_result.get(key)) and val not in (None, "", {})
        }

        if include_converted:
            obj_url = urls.get("obj")
            if obj_url:
                converted = self.resource_convert(
                    obj_url,
                    ["fbx", "stl", "usdz", "mp4", "gif"]
                )
                urls.update(converted)

        return urls

    def download_model(self, creation_id: str, format_key: str = "glb",
                       output_path: Optional[str] = None,
                       include_converted: bool = False) -> str:
        """
        下载指定格式的模型文件

        Args:
            creation_id: 创作ID
            format_key: 格式键名，如 'glb' / 'obj' / 'usdz' / 'pbrNormalImage' 等
            output_path: 本地保存路径，为空时自动从URL推断文件名
            include_converted: 是否包含转换格式（fbx/stl/usdz/mp4/gif）

        Returns:
            实际保存的本地文件路径
        """
        urls = self.get_model_urls(creation_id, include_converted=include_converted)
        if format_key not in urls:
            available = ", ".join(urls.keys())
            raise ValueError(
                f"格式 '{format_key}' 不可用。可用格式: {available}"
            )

        url = urls[format_key]
        resp = self.session.get(url, timeout=120)
        resp.raise_for_status()

        if output_path is None:
            parsed = urlparse(url)
            filename = unquote(os.path.basename(parsed.path)) or f"{creation_id}_{format_key}"
            output_path = filename

        with open(output_path, "wb") as f:
            f.write(resp.content)

        return os.path.abspath(output_path)

    def cancel_generation(self, creation_id: str) -> Dict[str, Any]:
        """取消生成任务"""
        return self._make_request("POST", "/creations/cancel", data={
            "creationsId": creation_id
        })

    # ========== 核心生成接口 ==========

    def _generate(self, model_type: str, scene_type: str = "playGround3D-2.0",
                  title: str = "", style: str = "", text: str = "",
                  image_list: Optional[List[str]] = None,
                  enable_pbr: bool = True, enable_low_poly: bool = False,
                  face_count: int = 1500000, motion_type: Optional[int] = None,
                  topology_format: Optional[str] = None) -> Dict[str, Any]:
        """通用生成接口"""
        data = {
            "sceneType": scene_type,
            "count": 1,
            "modelType": model_type,
            "title": title,
            "style": style,
            "enable_pbr": enable_pbr,
            "enableLowPoly": enable_low_poly,
            "faceCount": face_count
        }

        if text:
            data["prompt"] = text
        if image_list:
            data["imageList"] = image_list
        if motion_type is not None:
            data["motionType"] = motion_type
        if topology_format:
            data["topology_format"] = topology_format

        return self._make_request("POST", "/creations/generations", data=data)

    # ========== 图生3D ==========

    def generate_from_image(self, image_url: str, title: str = "",
                            style: str = "", enable_pbr: bool = True,
                            enable_low_poly: bool = False,
                            face_count: int = 1500000) -> Dict[str, Any]:
        """
        图生3D

        Args:
            image_url: 图片URL（需先上传获取resourceId）
            title: 作品标题
            style: 纹理风格
            enable_pbr: 是否启用PBR材质
            enable_low_poly: 是否低多边形
            face_count: 面数
        """
        return self._generate(
            model_type=self.MODEL_TYPE_IMAGE,
            title=title,
            style=style,
            image_list=[image_url],
            enable_pbr=enable_pbr,
            enable_low_poly=enable_low_poly,
            face_count=face_count
        )

    def generate_from_multi_view(self, image_urls: List[str], title: str = "",
                                 style: str = "", enable_pbr: bool = True,
                                 enable_low_poly: bool = False,
                                 face_count: int = 1500000) -> Dict[str, Any]:
        """
        多图视角生3D（多视图）

        Args:
            image_urls: 多角度图片URL列表（通常2-4张不同视角）
            title: 作品标题
            style: 纹理风格
            enable_pbr: 是否启用PBR材质
            enable_low_poly: 是否低多边形
            face_count: 面数
        """
        return self._generate(
            model_type=self.MODEL_TYPE_IMAGE,
            title=title,
            style=style,
            image_list=image_urls,
            enable_pbr=enable_pbr,
            enable_low_poly=enable_low_poly,
            face_count=face_count
        )

    # ========== 文生3D ==========

    def generate_from_text(self, prompt: str, title: str = "",
                           style: str = "", enable_pbr: bool = True,
                           enable_low_poly: bool = False,
                           face_count: int = 1500000, count: int = 4) -> Dict[str, Any]:
        """
        文生3D

        Args:
            prompt: 文本描述/提示词
            title: 作品标题
            style: 纹理风格
            enable_pbr: 是否启用PBR材质
            enable_low_poly: 是否低多边形
            face_count: 面数
            count: 生成数量，文生3D固定为4
        """
        data = {
            "sceneType": "playGround3D-2.0",
            "count": count,
            "modelType": self.MODEL_TYPE_TEXT,
            "title": title or prompt,
            "style": style,
            "prompt": prompt,
            "enable_pbr": enable_pbr,
            "enableLowPoly": enable_low_poly,
            "faceCount": face_count
        }
        return self._make_request("POST", "/creations/generations", data=data)

    # ========== 草图生3D ==========

    def generate_from_sketch(self, sketch_url: str, prompt: str,
                             title: str = "", style: str = "",
                             enable_pbr: bool = True,
                             enable_low_poly: bool = False,
                             face_count: int = 1500000) -> Dict[str, Any]:
        """
        草图生3D

        Args:
            sketch_url: 草图图片URL
            prompt: 草图描述提示词
            title: 作品标题
            style: 纹理风格
            enable_pbr: 是否启用PBR材质
            enable_low_poly: 是否低多边形
            face_count: 面数
        """
        return self._generate(
            model_type=self.MODEL_TYPE_SKETCH,
            title=title,
            style=style,
            text=prompt,
            image_list=[sketch_url],
            enable_pbr=enable_pbr,
            enable_low_poly=enable_low_poly,
            face_count=face_count
        )

    # ========== 3D动画生成 ==========

    def generate_animation(self, model_image_url: str, motion_type: int,
                           title: str = "") -> Dict[str, Any]:
        """
        3D动画生成

        Args:
            model_image_url: 3D模型图片URL
            motion_type: 动作类型ID
            title: 作品标题
        """
        return self._generate(
            model_type=self.MODEL_TYPE_ANIMATION,
            title=title,
            image_list=[model_image_url],
            motion_type=motion_type
        )

    # ========== 3D纹理生成 ==========

    def generate_texture(self, white_model_url: str, prompt: str,
                         title: str = "") -> Dict[str, Any]:
        """
        3D纹理生成

        Args:
            white_model_url: 白模图片URL
            prompt: 纹理描述
            title: 作品标题
        """
        return self._generate(
            model_type=self.MODEL_TYPE_TEXTURE,
            title=title,
            text=prompt,
            image_list=[white_model_url]
        )

    # ========== 3D智能拓扑 ==========

    def generate_lowpoly(self, model_url: str, face_count: int = 5000,
                         topology_format: str = "glb",
                         title: str = "") -> Dict[str, Any]:
        """
        3D智能拓扑（减面）

        Args:
            model_url: 模型图片URL
            face_count: 目标面数（5000/18000/30000）
            topology_format: 输出格式（glb/obj）
            title: 作品标题
        """
        return self._generate(
            model_type=self.MODEL_TYPE_LOWPOLY,
            title=title,
            image_list=[model_url],
            enable_low_poly=True,
            face_count=face_count,
            topology_format=topology_format
        )

    # ========== 资源上传 ==========

    def get_upload_info(self, filename: str) -> Dict[str, Any]:
        """获取上传凭证"""
        return self._make_request("POST", "/resource/genUploadInfo", data={
            "fileName": filename
        })

    def review_resource(self, resource_url: str, scene_type: str = "playGround3D-2.0",
                        resource_type: str = "image", text: str = "") -> Dict[str, Any]:
        """资源审核"""
        return self._make_request("POST", "/resource/review", data={
            "sceneType": scene_type,
            "text": text,
            "resourceType": resource_type,
            "resourceUrl": resource_url
        })

    # ========== 分享 ==========

    def create_share(self, creation_id: str, platform: str = "3dPlayground") -> Dict[str, Any]:
        """创建分享"""
        return self._make_request("POST", "/share", data={
            "contentType": "creation",
            "contentId": creation_id,
            "sharedContent": "",
            "platform": platform
        })

    # ========== 配置 ==========

    def get_config(self) -> Dict[str, Any]:
        """获取全局配置"""
        return self._make_request("GET", "/config")

    def get_action_templates(self) -> Dict[str, Any]:
        """获取动画动作模板"""
        return self._make_request("GET", "/workflow/action/templates")

    # ========== 轮询等待 ==========

    def wait_for_completion(self, creation_id: str, timeout: int = 600,
                           poll_interval: int = 5) -> Dict[str, Any]:
        """
        等待生成完成

        Args:
            creation_id: 创作ID
            timeout: 超时时间（秒）
            poll_interval: 轮询间隔（秒）
        """
        start_time = time.time()
        while time.time() - start_time < timeout:
            status = self.get_generation_status(creation_id)

            # Handle both wrapped {code, data} and flat response formats
            if "code" in status:
                data = status.get("data", {})
            else:
                data = status
            state = data.get("status")

            if state == "success":
                return status
            elif state == "fail":
                raise Exception(f"Generation failed: {data}")

            progress = data.get("progress", 0)
            print(f"State: {state}, progress: {progress}%")

            time.sleep(poll_interval)

        raise TimeoutError(f"Generation timeout after {timeout} seconds")


from .config import get_cookie_path

def load_cookies_from_file(filepath: Optional[str] = None) -> str:
    """从文件加载Cookie，默认读取用户配置目录的 cookies.txt"""
    path = filepath or str(get_cookie_path())
    with open(path, 'r', encoding='utf-8') as f:
        return f.read().strip()


if __name__ == "__main__":
    import sys

    cookies = load_cookies_from_file(sys.argv[1] if len(sys.argv) > 1 else None)
    api = Hunyuan3DAPI(cookies)

    # 测试配额查询
    print("查询配额...")
    quota = api.get_quota_info()
    print(json.dumps(quota, indent=2, ensure_ascii=False))

    # 测试获取配置
    print("\n获取配置...")
    config = api.get_config()
    print(f"Config keys: {list(config.keys())}")