#!/usr/bin/env python3
"""
腾讯混元3D纯Python API客户端
无需浏览器，直接HTTP调用
"""

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

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


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

    # 模型下载格式键名
    # 原生格式：直接来自 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):
        """
        初始化API客户端

        Args:
            cookies: 浏览器Cookie字符串，包含hunyuan_user和hunyuan_token
        """
        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):
        """设置Cookie"""
        self.session.headers["Cookie"] = cookies

    def set_cookie_dict(self, cookie_dict: Dict[str, str]):
        """从字典设置Cookie"""
        cookie_str = "; ".join(f"{k}={v}" for k, v in cookie_dict.items())
        self.session.headers["Cookie"] = cookie_str

    def _make_request(self, method: str, endpoint: str, params: Optional[Dict] = None,
                      data: Optional[Dict] = None) -> Dict[str, Any]:
        """
        发送带签名的API请求

        注意: 腾讯混元3D的签名只包含URL查询参数，不包含请求体数据
        """
        url = f"{API_BASE}{endpoint}"

        # 签名只针对查询参数，不包含请求体
        sign_params = {}
        if params:
            sign_params.update(params)

        # 生成签名
        signed_params = sign(sign_params)

        # 分离签名参数
        timestamp = signed_params.pop("timestamp")
        nonce = signed_params.pop("nonce")
        sign_value = signed_params.pop("sign")

        # 构建最终URL查询参数（只包含签名相关）
        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:
            # POST请求：签名在URL参数中，数据在body中
            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 generate_3d(self, image_url: str, scene_type: str = "playGround3D-2.0",
                    model_type: str = "image2ModelV3.1", title: str = "",
                    enable_pbr: bool = True, enable_low_poly: bool = False,
                    face_count: int = 1500000, count: int = 1) -> Dict[str, Any]:
        """
        生成3D模型

        Args:
            image_url: 图片URL（需要先上传到腾讯云COS）
            scene_type: 场景类型
            model_type: 模型类型
            title: 标题
            enable_pbr: 是否启用PBR材质
            enable_low_poly: 是否启用低多边形
            face_count: 面数
            count: 生成数量

        Returns:
            包含creationsId的响应
        """
        data = {
            "sceneType": scene_type,
            "count": count,
            "modelType": model_type,
            "title": title,
            "style": "",
            "imageList": [image_url],
            "enable_pbr": enable_pbr,
            "enableLowPoly": enable_low_poly,
            "faceCount": face_count
        }

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

    def generate_text(self, prompt: str, scene_type: str = "playGround3D-2.0",
                      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: 文本描述/提示词
            scene_type: 场景类型
            title: 作品标题
            style: 纹理风格
            enable_pbr: 是否启用PBR材质
            enable_low_poly: 是否低多边形
            face_count: 面数
            count: 生成数量，文生3D固定为4

        Returns:
            包含creationsId的响应
        """
        data = {
            "sceneType": scene_type,
            "count": count,
            "modelType": "text2ModelV3.1",
            "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)

    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' / 'pbrNormalImage' 等
            output_path: 本地保存路径，为空时自动从URL推断文件名

        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 wait_for_completion(self, creation_id: str, timeout: int = 600,
                           poll_interval: int = 5) -> Dict[str, Any]:
        """
        等待生成完成

        Args:
            creation_id: 创作ID
            timeout: 超时时间（秒）
            poll_interval: 轮询间隔（秒）

        Returns:
            完成的创作信息
        """
        start_time = time.time()
        last_error = None
        while time.time() - start_time < timeout:
            try:
                status = self.get_generation_status(creation_id)
            except Exception as e:
                last_error = e
                print(f"查询状态失败: {e}，将在 {poll_interval} 秒后重试...")
                time.sleep(poll_interval)
                continue

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

            if state == "success":
                return status
            elif state == "fail":
                raise Exception(f"Generation failed: {data.get('errorMsg', 'Unknown error')}")

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

            time.sleep(poll_interval)

        if last_error:
            raise TimeoutError(f"Generation timeout after {timeout} seconds. Last error: {last_error}")
        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

    if len(sys.argv) > 1 and sys.argv[1] == "--cookie-string":
        cookies = sys.argv[2]
    else:
        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))