#!/usr/bin/env python3
"""Meshy text-to-3d CLI（Python 3 标准库，无需额外依赖）。

用法：
    MESHY_API_KEY=msy_xxx python3 meshy_gen.py "<prompt>" "<输出.glb>" [texture_prompt]

流程：preview -> 轮询 -> refine(PBR) -> 轮询 -> 下载 GLB。
"""
import argparse
import json
import os
import shutil
import sys
import time
import urllib.error
import urllib.request

API_URL = "https://api.meshy.ai/openapi/v2/text-to-3d"


def load_key():
    key = os.environ.get("MESHY_API_KEY", "").strip()
    if key:
        return key
    config_path = os.path.expanduser("~/.config/meshy_api_key")
    if os.path.exists(config_path):
        key = open(config_path).read().strip()
        if key:
            return key
    sys.exit("没有 Meshy API Key。设置环境变量 MESHY_API_KEY，或写入 ~/.config/meshy_api_key")


def request(method, url, body=None, key=None):
    req = urllib.request.Request(url, method=method)
    req.add_header("Authorization", "Bearer " + key)
    req.add_header("Content-Type", "application/json")
    data = json.dumps(body).encode() if body is not None else None
    try:
        with urllib.request.urlopen(req, data=data) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        text = e.read().decode()
        try:
            text = json.loads(text).get("message", text)
        except json.JSONDecodeError:
            pass
        sys.exit(f"HTTP {e.code}: {text}")


def create_task(mode, prompt, key, preview_task_id=None, texture_prompt=None):
    body = {
        "mode": mode,
        "prompt": prompt,
        "model_type": "lowpoly",
        "target_formats": ["glb"],
        "pose_mode": "",
        "should_remesh": False,
        "art_style": "realistic",
    }
    if mode == "refine":
        body["preview_task_id"] = preview_task_id
        body["enable_pbr"] = True
        body["remove_lighting"] = True
        if texture_prompt:
            body["texture_prompt"] = texture_prompt
    return request("POST", API_URL, body, key)["result"]


def poll_task(task_id, key):
    url = f"{API_URL}/{task_id}"
    while True:
        result = request("GET", url, None, key)
        status = result.get("status")
        if status == "SUCCEEDED":
            return result
        if status == "FAILED":
            message = result.get("task_error", {}).get("message", "unknown error")
            sys.exit(f"任务失败：{message}")
        print(f"  {status} ...", flush=True)
        time.sleep(5)


def download(url, output):
    req = urllib.request.Request(url)
    with urllib.request.urlopen(req) as resp, open(output, "wb") as f:
        shutil.copyfileobj(resp, f)


def main():
    parser = argparse.ArgumentParser(description="Meshy text-to-3d 生成 GLB")
    parser.add_argument("prompt", help="模型描述，必须包含 stylized, low-poly, game asset")
    parser.add_argument("output", help="输出 .glb 路径")
    parser.add_argument("texture_prompt", nargs="?", default="", help="可选贴图提示")
    args = parser.parse_args()

    key = load_key()
    print("创建 preview 任务 ...", flush=True)
    preview_id = create_task("preview", args.prompt, key)
    print(f"  preview task: {preview_id}", flush=True)
    preview = poll_task(preview_id, key)
    print("  preview 完成", flush=True)

    print("创建 refine 任务（PBR 贴图）...", flush=True)
    refine_id = create_task(
        "refine",
        args.prompt,
        key,
        preview_task_id=preview_id,
        texture_prompt=args.texture_prompt,
    )
    print(f"  refine task: {refine_id}", flush=True)
    refine = poll_task(refine_id, key)
    print("  refine 完成", flush=True)

    glb_url = refine.get("model_urls", {}).get("glb")
    if not glb_url:
        sys.exit("响应中没有 model_urls.glb")
    download(glb_url, args.output)
    print(f"已保存: {args.output}")


if __name__ == "__main__":
    main()