#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import re
import sys
import json
import time
import argparse
import mimetypes
from pathlib import Path
from typing import Any, Dict, Iterator, Tuple
from urllib.parse import urlparse

import requests

BASE_URL = "https://api.tripo3d.ai/v2/openapi"


class TripoAPIError(RuntimeError):
    pass


def build_parser():
    p = argparse.ArgumentParser("Tripo3D CLI: single image -> 3D")

    # I/O
    p.add_argument("-i", "--image", required=True, help="Input image path")
    p.add_argument("-o", "--out_dir", default="tripo_outputs", help="Output directory")

    # Auth
    p.add_argument(
        "--api_key",
        default=os.getenv("TRIPO_API_KEY", "tcli_50ecbff125084d4db958b1863ec082e6"),
        help="Tripo API key, or set env TRIPO_API_KEY",
    )

    # Model
    p.add_argument(
        "--model_version",
        type=str,
        default="v3.1-20260211",
        help="Model version, e.g. P1-20260311 / v3.1-20260211 / v3.0-20250812 / v2.5-20250123",
    )

    # Network / polling
    p.add_argument("--poll_interval", type=float, default=2.0, help="Polling interval (seconds)")
    p.add_argument("--poll_timeout", type=float, default=1800.0, help="Max polling time (seconds)")
    p.add_argument("--request_timeout", type=float, default=120.0, help="HTTP request timeout (seconds)")

    # Generation options
    p.add_argument("--texture", dest="texture", action="store_true", default=True)
    p.add_argument("--no-texture", dest="texture", action="store_false")

    p.add_argument("--pbr", dest="pbr", action="store_true", default=True)
    p.add_argument("--no-pbr", dest="pbr", action="store_false")

    p.add_argument(
        "--texture_quality",
        type=str,
        default="standard",
        choices=["standard", "detailed"],
        help="Texture quality",
    )
    p.add_argument(
        "--texture_alignment",
        type=str,
        default="original_image",
        choices=["original_image", "geometry"],
        help="Texture alignment mode",
    )
    p.add_argument(
        "--orientation",
        type=str,
        default="default",
        choices=["default", "align_image"],
        help="Orientation mode",
    )

    # Optional params
    p.add_argument("--face_limit", type=int, default=None)
    p.add_argument("--model_seed", type=int, default=None)
    p.add_argument("--texture_seed", type=int, default=None)
    p.add_argument("--auto_size", type=str, default=None)
    p.add_argument("--quad", type=str, default=None)
    p.add_argument("--compress", type=str, default=None)
    p.add_argument("--generate_parts", type=str, default=None)
    p.add_argument("--smart_low_poly", type=str, default=None)

    # Save / download toggles
    p.add_argument("--download_outputs", dest="download_outputs", action="store_true", default=True)
    p.add_argument("--no-download_outputs", dest="download_outputs", action="store_false")

    p.add_argument("--save_task_json", dest="save_task_json", action="store_true", default=True)
    p.add_argument("--no-save_task_json", dest="save_task_json", action="store_false")

    p.add_argument("--print_payload", dest="print_payload", action="store_true", default=False)
    p.add_argument("--print_output", dest="print_output", action="store_true", default=True)
    p.add_argument("--no-print_output", dest="print_output", action="store_false")

    return p


def guess_mime_type(file_path: Path) -> str:
    mime, _ = mimetypes.guess_type(str(file_path))
    return mime or "application/octet-stream"


def safe_filename(name: str) -> str:
    name = re.sub(r'[\\/:*?"<>|]+', "_", name)
    name = re.sub(r"\s+", "_", name).strip("._")
    return name or "file"


def extract_error_message(payload: Any) -> str:
    if isinstance(payload, dict):
        for key in ("message", "error", "error_message", "detail"):
            if payload.get(key):
                return str(payload[key])

        data = payload.get("data")
        if isinstance(data, dict):
            for key in ("message", "error", "error_message", "detail"):
                if data.get(key):
                    return str(data[key])

        return json.dumps(payload, ensure_ascii=False)[:800]

    return str(payload)[:800]


def request_json(
        session: requests.Session,
        method: str,
        endpoint: str,
        request_timeout: float,
        **kwargs,
) -> Dict[str, Any]:
    url = f"{BASE_URL}{endpoint}"

    try:
        resp = session.request(method=method, url=url, timeout=request_timeout, **kwargs)
    except requests.RequestException as e:
        raise TripoAPIError(f"请求失败: {method} {url} | {e}") from e

    if not resp.ok:
        try:
            err_payload = resp.json()
        except Exception:
            err_payload = resp.text
        raise TripoAPIError(
            f"HTTP {resp.status_code} | {method} {url} | {extract_error_message(err_payload)}"
        )

    try:
        payload = resp.json()
    except Exception as e:
        raise TripoAPIError(
            f"响应不是合法 JSON: {method} {url}\n原始响应前 500 字符:\n{resp.text[:500]}"
        ) from e

    return payload


def create_session(api_key: str) -> requests.Session:
    session = requests.Session()
    session.headers.update({
        "Authorization": f"Bearer {api_key}",
        "Accept": "application/json",
    })
    return session


def upload_image(session: requests.Session, image_path: Path, request_timeout: float) -> str:
    if not image_path.exists():
        raise FileNotFoundError(f"找不到图片: {image_path}")

    with image_path.open("rb") as f:
        files = {
            "file": (image_path.name, f, guess_mime_type(image_path))
        }
        payload = request_json(
            session,
            "POST",
            "/upload",
            request_timeout=request_timeout,
            files=files,
        )

    data = payload.get("data") or {}
    file_token = data.get("image_token")

    if not file_token:
        raise TripoAPIError(f"上传成功但未返回 image_token: {json.dumps(payload, ensure_ascii=False)}")

    return file_token


def build_generation_payload(args, file_token: str, image_path: Path) -> Dict[str, Any]:
    file_ext = image_path.suffix.lower().lstrip(".") or "png"

    payload: Dict[str, Any] = {
        "type": "image_to_model",
        "model_version": args.model_version,
        "file": {
            "type": file_ext,
            "file_token": file_token,
        },
        "texture": args.texture,
        "pbr": args.pbr,
        "texture_quality": args.texture_quality,
        "texture_alignment": args.texture_alignment,
        "orientation": args.orientation,
    }

    optional_fields = [
        "face_limit",
        "model_seed",
        "texture_seed",
        "auto_size",
        "quad",
        "compress",
        "generate_parts",
        "smart_low_poly",
    ]

    for key in optional_fields:
        value = getattr(args, key)
        if value is not None:
            payload[key] = value

    return payload


def create_task(session: requests.Session, payload: Dict[str, Any], request_timeout: float) -> str:
    resp = request_json(
        session,
        "POST",
        "/task",
        request_timeout=request_timeout,
        json=payload,
    )
    data = resp.get("data") or {}
    task_id = data.get("task_id")

    if not task_id:
        raise TripoAPIError(f"提交任务成功但未返回 task_id: {json.dumps(resp, ensure_ascii=False)}")

    return task_id


def get_task(session: requests.Session, task_id: str, request_timeout: float) -> Dict[str, Any]:
    return request_json(
        session,
        "GET",
        f"/task/{task_id}",
        request_timeout=request_timeout,
    )


def poll_task(
        session: requests.Session,
        task_id: str,
        poll_interval: float,
        poll_timeout: float,
        request_timeout: float,
) -> Dict[str, Any]:
    start = time.perf_counter()
    last_line = ""

    while True:
        resp = get_task(session, task_id, request_timeout=request_timeout)
        data = resp.get("data") or {}

        status = str(data.get("status", "unknown")).lower()
        progress = data.get("progress", 0)
        elapsed = time.perf_counter() - start

        line = f"\r[状态] {status:<10} | [进度] {progress:>3}% | [已等待] {elapsed:>7.1f}s"
        if line != last_line:
            sys.stdout.write(line)
            sys.stdout.flush()
            last_line = line

        if status == "success":
            sys.stdout.write("\n")
            sys.stdout.flush()
            return resp

        if status == "failed":
            sys.stdout.write("\n")
            sys.stdout.flush()
            error_message = data.get("error_message") or extract_error_message(resp)
            raise TripoAPIError(f"任务失败 | task_id={task_id} | {error_message}")

        if elapsed > poll_timeout:
            sys.stdout.write("\n")
            sys.stdout.flush()
            raise TimeoutError(f"轮询超时: 已等待 {elapsed:.1f}s，task_id={task_id}")

        time.sleep(poll_interval)


def iter_urls(obj: Any, prefix: str = "output") -> Iterator[Tuple[str, str]]:
    if isinstance(obj, dict):
        for k, v in obj.items():
            yield from iter_urls(v, f"{prefix}.{k}")
    elif isinstance(obj, list):
        for i, v in enumerate(obj):
            yield from iter_urls(v, f"{prefix}[{i}]")
    elif isinstance(obj, str) and obj.startswith(("http://", "https://")):
        yield prefix, obj


def infer_extension_from_url(url: str) -> str:
    path = urlparse(url).path
    ext = Path(path).suffix
    return ext if ext else ".bin"


def unique_path(path: Path) -> Path:
    if not path.exists():
        return path

    stem = path.stem
    suffix = path.suffix
    parent = path.parent
    i = 1
    while True:
        candidate = parent / f"{stem}_{i}{suffix}"
        if not candidate.exists():
            return candidate
        i += 1


def download_file(session: requests.Session, url: str, save_path: Path, request_timeout: float) -> None:
    try:
        with session.get(url, stream=True, timeout=request_timeout) as resp:
            resp.raise_for_status()
            with save_path.open("wb") as f:
                for chunk in resp.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        f.write(chunk)
    except requests.RequestException as e:
        raise TripoAPIError(f"下载失败: {url} | {e}") from e


def save_outputs(
        session: requests.Session,
        task_resp: Dict[str, Any],
        out_dir: Path,
        request_timeout: float,
        save_task_json: bool = True,
        download_outputs: bool = True,
) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)

    data = task_resp.get("data") or {}
    task_id = data.get("task_id", "unknown_task")
    output = data.get("output") or {}

    if save_task_json:
        meta_path = out_dir / f"{safe_filename(task_id)}.json"
        with meta_path.open("w", encoding="utf-8") as f:
            json.dump(task_resp, f, ensure_ascii=False, indent=2)

    if not output:
        print("⚠️ 任务成功，但 output 为空。")
        return

    if not download_outputs:
        print("ℹ️ 已跳过下载，仅保存任务响应。")
        return

    url_items = list(iter_urls(output))
    if not url_items:
        print("⚠️ output 中没有找到可下载 URL。")
        return

    print("\n📥 开始下载输出文件...")
    for logical_key, url in url_items:
        short_key = logical_key.replace("output.", "")
        ext = infer_extension_from_url(url)
        filename = safe_filename(short_key) + ext
        save_path = unique_path(out_dir / filename)

        print(f"  - {short_key} -> {save_path}")
        download_file(session, url, save_path, request_timeout=request_timeout)


def main():
    args = build_parser().parse_args()

    if not args.api_key:
        raise ValueError("请提供 --api_key 或设置环境变量 TRIPO_API_KEY")

    image_path = Path(args.image)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    session = create_session(args.api_key)

    print(f"🚀 启动测试 | 模型: {args.model_version}")
    print(f"🖼️ 输入图片: {image_path}")
    print(f"📁 输出目录: {out_dir.resolve()}")

    start_wall_time = time.perf_counter()

    # 1) 上传
    print("\n[1/4] 上传图片...")
    upload_start = time.perf_counter()
    file_token = upload_image(session, image_path, request_timeout=args.request_timeout)
    upload_end = time.perf_counter()
    print(f"✅ 上传完成 | file_token: {file_token}")
    print(f"⏱️ 上传耗时: {upload_end - upload_start:.2f}s")

    # 2) 提交任务
    print("\n[2/4] 提交 image_to_model 任务...")
    payload = build_generation_payload(args, file_token, image_path)

    if args.print_payload:
        print(json.dumps(payload, ensure_ascii=False, indent=2))

    task_id = create_task(session, payload, request_timeout=args.request_timeout)
    print(f"✅ 任务提交成功 | task_id: {task_id}")

    # 3) 轮询任务
    print("\n[3/4] 轮询任务状态...")
    gen_start = time.perf_counter()
    task_resp = poll_task(
        session,
        task_id,
        poll_interval=args.poll_interval,
        poll_timeout=args.poll_timeout,
        request_timeout=args.request_timeout,
    )
    gen_end = time.perf_counter()

    data = task_resp.get("data") or {}
    output = data.get("output") or {}
    total_end = time.perf_counter()

    print("\n🎉 生成成功")
    print("=" * 60)
    print(f"task_id             : {task_id}")
    print(f"纯生成耗时          : {gen_end - gen_start:.2f}s")
    print(f"总流程耗时          : {total_end - start_wall_time:.2f}s")
    print(f"最终 status         : {data.get('status')}")
    print(f"output keys         : {list(output.keys()) if isinstance(output, dict) else type(output)}")
    print("=" * 60)

    if args.print_output:
        print(json.dumps(output, ensure_ascii=False, indent=2))

    # 4) 下载输出
    print("\n[4/4] 保存结果...")
    save_outputs(
        session,
        task_resp,
        out_dir=out_dir,
        request_timeout=args.request_timeout,
        save_task_json=args.save_task_json,
        download_outputs=args.download_outputs,
    )

    print("\n✅ 全部完成。")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n❌ 程序终止: {e}")
        sys.exit(1)