#!/usr/bin/env python3
"""非人角色：已验收游戏模 → TokenRig 绑骨 → UniMate 三动作（Idle/Walk/Attack）。

正确流程（与 GameResProductProc 二.3.B 一致）：
  1. 原混元网页 + model_lowpoly_bake + SetupCharacter 做出静态模（XPbr）
  2. 本脚本只吃已验收网格（FBX/GLB），不再走 DGX Hunyuan3D 重做模
  3. TokenRig + UniMate（经本机 LLM3DGraphics backend，SSH 到 DGXA）

示例：
  python Tools/dgx3d_animal_pipeline.py \\
    --mesh AIClient/Assets/Bundles/Characters/Animal/RockBall/RockBall.fbx \\
    -o Tools/output/rockball_dgx \\
    --name RockBall
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
import uuid
import zipfile
from pathlib import Path
from typing import Dict, List, Optional, Tuple

REPO = Path(__file__).resolve().parents[1]
DEFAULT_LLM3D = Path(
    os.environ.get("LLM3D_ROOT", r"D:\UD\AI\XModel\LLM3DGraphics")
)

DEFAULT_MOTIONS = [
    ("Idle", "a small bipedal character idles with a gentle breathing motion"),
    ("Walk", "a small bipedal character walks forward upright with alternating short steps and gentle arm swings"),
    ("Attack", "a small bipedal character makes a short forward body strike"),
]

_BLENDER_CONVERT = r'''
import sys
from pathlib import Path
import bpy
from mathutils import Matrix
from math import pi

argv = sys.argv[sys.argv.index("--") + 1:]
src, dst = Path(argv[0]), Path(argv[1])
bpy.ops.wm.read_factory_settings(use_empty=True)
ext = src.suffix.lower()
if ext == ".fbx":
    bpy.ops.import_scene.fbx(filepath=str(src))
    # Project static FBX stores Y-up vertices. TokenRig/glTF requires Blender
    # Z-up before the glTF exporter performs its Y-up conversion.
    if any(o.type == 'ARMATURE' for o in bpy.context.scene.objects):
        raise SystemExit('Expected approved STATIC mesh, not an already skinned FBX')
    for obj in bpy.context.scene.objects:
        if obj.parent is None:
            obj.matrix_world = Matrix.Rotation(pi / 2, 4, 'X') @ obj.matrix_world
elif ext in (".glb", ".gltf"):
    bpy.ops.import_scene.gltf(filepath=str(src))
else:
    raise SystemExit(f"unsupported: {ext}")
# Prefer sibling albedo if FBX came without embeds
parent = src.parent
for name in (src.stem + "_base.png", src.stem + "_color.png", "diffuse.png"):
    tex = parent / name
    if not tex.is_file():
        continue
    img = bpy.data.images.load(str(tex), check_existing=True)
    for mat in bpy.data.materials:
        mat.use_nodes = True
        nt = mat.node_tree
        nodes, links = nt.nodes, nt.links
        principled = next((n for n in nodes if n.type == "BSDF_PRINCIPLED"), None)
        if principled is None:
            continue
        tex_node = nodes.new("ShaderNodeTexImage")
        tex_node.image = img
        links.new(tex_node.outputs["Color"], principled.inputs["Base Color"])
    break
dst.parent.mkdir(parents=True, exist_ok=True)
bpy.ops.export_scene.gltf(filepath=str(dst), export_format="GLB", export_yup=True)
print("OK", dst)
'''


def _ensure_llm3d() -> Path:
    root = DEFAULT_LLM3D.resolve()
    if not (root / "backend" / "tokenrig_runner.py").is_file():
        raise SystemExit(f"LLM3DGraphics 未找到: {root}（可用 LLM3D_ROOT 覆盖）")
    if str(root) not in sys.path:
        sys.path.insert(0, str(root))
    return root


def _blender() -> Path:
    _ensure_llm3d()
    from backend.settings import blender_path

    path = blender_path()
    if not path:
        raise SystemExit("未配置 Blender（LLM3DGraphics blender_path）")
    return Path(path)


def require_dgx_backends() -> None:
    _ensure_llm3d()
    from backend.tokenrig_runner import tokenrig_ready
    from backend.unimate_runner import unimate_ready

    tr, uni = bool(tokenrig_ready()), bool(unimate_ready())
    if not tr or not uni:
        raise SystemExit(
            f"DGX 能力红灯：tokenrig_ready={tr} unimate_ready={uni}\n"
            "请先 nvsync connect DGXA --detach，并确认 SkinTokens/UniMate 就绪。"
        )


def mesh_to_glb(mesh: Path, out_glb: Path) -> Path:
    mesh = mesh.resolve()
    out_glb = out_glb.resolve()
    out_glb.parent.mkdir(parents=True, exist_ok=True)
    if mesh.suffix.lower() == ".glb":
        if mesh.resolve() != out_glb:
            shutil.copy2(mesh, out_glb)
        return out_glb
    script = out_glb.parent / "_fbx_to_glb.py"
    script.write_text(_BLENDER_CONVERT, encoding="utf-8")
    blender = _blender()
    proc = subprocess.run(
        [
            str(blender),
            "--background",
            "--factory-startup",
            "--python",
            str(script),
            "--",
            str(mesh),
            str(out_glb),
        ],
        capture_output=True,
        text=True,
        timeout=300,
        check=False,
    )
    if proc.returncode != 0 or not out_glb.is_file():
        raise SystemExit(
            f"FBX→GLB 失败 rc={proc.returncode}\n"
            f"{(proc.stdout or '')[-1500:]}\n{(proc.stderr or '')[-1500:]}"
        )
    print(f"  mesh glb={out_glb} ({out_glb.stat().st_size} bytes)", flush=True)
    return out_glb


def extract_rigged_fbx(zip_or_fbx: Path, out_dir: Path) -> Path:
    out_dir.mkdir(parents=True, exist_ok=True)
    if zip_or_fbx.suffix.lower() == ".fbx":
        dest = out_dir / zip_or_fbx.name
        if zip_or_fbx.resolve() != dest.resolve():
            shutil.copy2(zip_or_fbx, dest)
        return dest
    with zipfile.ZipFile(zip_or_fbx, "r") as zf:
        zf.extractall(out_dir)
    candidates = list(out_dir.rglob("model.fbx")) + list(out_dir.rglob("*.fbx"))
    if not candidates:
        raise FileNotFoundError(f"no fbx in {zip_or_fbx}")
    for c in candidates:
        if c.name.lower() == "model.fbx":
            return c
    return candidates[0]


def parse_motions(raw: Optional[str], name: str) -> List[Tuple[str, str]]:
    if not raw:
        if name.lower() != "rockball":
            return [
                ("Idle", f"{name}宠物原地待机，轻微呼吸，可循环"),
                ("Walk", f"{name}宠物原地走路，可循环"),
                ("Attack", f"{name}宠物原地短促攻击"),
            ]
        return list(DEFAULT_MOTIONS)
    items = json.loads(raw)
    return [(str(it["name"]), str(it["prompt"])) for it in items]


def run_pipeline(
    mesh: Path,
    out_dir: Path,
    *,
    name: str,
    motions: List[Tuple[str, str]],
    skip_gate: bool = False,
) -> dict:
    mesh = mesh.resolve()
    if not mesh.is_file():
        raise SystemExit(f"mesh not found: {mesh}")
    out_dir = out_dir.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    if not skip_gate:
        require_dgx_backends()
    else:
        print("WARNING: --skip-gate", flush=True)

    _ensure_llm3d()
    from backend.motion_branches import export_rigged_fbx
    from backend.tokenrig_runner import run_tokenrig
    import backend.unimate_runner as unimate_runner
    from backend.unimate_runner import submit_unimate_job

    # Preserve explicit English captions: upstream rewrites every English walk
    # to "a dog" and gives Chinese 跳 priority over 走. Log submitted prompts.
    original_normalize = unimate_runner.normalize_unimate_prompt
    def normalize_prompt(prompt):
        if not any('\u4e00' <= char <= '\u9fff' for char in prompt):
            return prompt.strip()
        return original_normalize(prompt)
    unimate_runner.normalize_unimate_prompt = normalize_prompt

    # DGXA 上 HF 偶发超时：强制离线用本地 flan-t5 缓存
    _orig_ssh = unimate_runner._ssh

    def _ssh_offline(host, cmd, timeout=120):
        if "run_job.sh" in str(cmd):
            cmd = (
                "export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 "
                "HF_DATASETS_OFFLINE=1; " + cmd
            )
        return _orig_ssh(host, cmd, timeout=timeout)

    unimate_runner._ssh = _ssh_offline

    job_id = uuid.uuid4().hex
    job_dir = out_dir / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    log_path = job_dir / "pipeline.log"
    print(f"job {job_id} mesh={mesh}", flush=True)

    mesh_glb = mesh_to_glb(mesh, job_dir / "model_ex.glb")
    print("TokenRig …", flush=True)
    rig_path = run_tokenrig(
        mesh_glb, job_dir, log_path, skeleton_template="original"
    )
    print(f"  rigged={rig_path}", flush=True)

    zip_cand = job_dir / "rigged_fbx.zip"
    print("export rigged FBX zip …", flush=True)
    exported = export_rigged_fbx(rig_path, zip_cand, log_path)
    if exported is None or not zip_cand.is_file():
        raise SystemExit("rigged FBX 导出失败，见 pipeline.log")

    rig_extract = job_dir / "rigged"
    model_fbx = extract_rigged_fbx(zip_cand, rig_extract)
    named = job_dir / f"{name}_Rigged.fbx"
    shutil.copy2(model_fbx, named)

    anim_dir = job_dir / "anims"
    anim_dir.mkdir(exist_ok=True)
    anim_paths: Dict[str, str] = {}
    for clip_name, prompt in motions:
        print(f"animate {clip_name}: {prompt}", flush=True)
        stamp = f"{job_id}_{clip_name.lower()}"
        last_err: Optional[BaseException] = None
        for attempt in range(1, 4):
            try:
                submit_unimate_job(
                    rig_path=rig_path,
                    prompt=prompt,
                    directory=job_dir,
                    log_path=log_path,
                    job_id=f"{stamp}_a{attempt}",
                )
                last_err = None
                break
            except Exception as exc:  # noqa: BLE001
                last_err = exc
                print(f"  UniMate {clip_name} attempt {attempt} failed: {exc}", flush=True)
                time.sleep(5)
        if last_err is not None:
            raise SystemExit(f"UniMate {clip_name} failed: {last_err}") from last_err
        src_fbx = job_dir / "animated.fbx"
        if not src_fbx.is_file():
            raise SystemExit(f"UniMate 未产出 animated.fbx ({clip_name})")
        dest = anim_dir / f"{name}_{clip_name}.fbx"
        shutil.copy2(src_fbx, dest)
        anim_paths[clip_name] = str(dest.relative_to(out_dir))
        glb = job_dir / "animated.glb"
        if glb.is_file():
            shutil.copy2(glb, dest.with_suffix(".glb"))
        print(f"  saved {dest}", flush=True)

    manifest = {
        "ok": True,
        "name": name,
        "job_id": job_id,
        "source_mesh": str(mesh),
        "pipeline": "approved_mesh→TokenRig→UniMate",
        "rigged_fbx": str(named),
        "rigged_zip": str(zip_cand) if zip_cand.is_file() else str(exported),
        "anims": anim_paths,
        "motions": [{"name": n, "prompt": p} for n, p in motions],
        "hint": "Unity: SetupAnimalAnimator 挂回静态模 XPbr；动画用 Anims/*.fbx",
    }
    (out_dir / f"{name}_manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    (job_dir / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(json.dumps(manifest, ensure_ascii=False, indent=2), flush=True)
    return manifest


def main(argv: Optional[List[str]] = None) -> int:
    ap = argparse.ArgumentParser(
        description="已验收网格 → TokenRig + UniMate（不重做 DGX 模型）"
    )
    ap.add_argument(
        "--mesh",
        "-m",
        type=Path,
        required=False,
        help="已验收游戏模 FBX/GLB（推荐 Bundles/.../{Name}.fbx）",
    )
    ap.add_argument("-o", "--out", type=Path, required=False)
    ap.add_argument("--name", default=None, help="English asset name, e.g. RockBall")
    ap.add_argument("--motions-json", default=None)
    ap.add_argument("--skip-gate", action="store_true")
    ap.add_argument(
        "--caps-only",
        action="store_true",
        help="只检查 TokenRig/UniMate 就绪",
    )
    ap.add_argument(
        "-i",
        "--image",
        type=Path,
        default=None,
        help=argparse.SUPPRESS,  # 旧图生入口已废弃
    )
    args = ap.parse_args(argv)

    if args.caps_only:
        require_dgx_backends()
        print(json.dumps({"tokenrig_ready": True, "unimate_ready": True}, indent=2))
        return 0

    if args.image is not None:
        raise SystemExit(
            "已废弃：勿再用 -i 设计图走 DGX Hunyuan3D 重做模。\n"
            "请先按原流程烘焙验收静态模，再：\n"
            "  python Tools/dgx3d_animal_pipeline.py --mesh "
            "AIClient/Assets/Bundles/Characters/Animal/{Name}/{Name}.fbx "
            "-o Tools/output/{name}_dgx --name {Name}"
        )

    if not args.mesh or not args.out or not args.name:
        raise SystemExit("需要 --mesh / -o / --name（或 --caps-only）")

    motions = parse_motions(args.motions_json, args.name)
    t0 = time.time()
    run_pipeline(
        args.mesh,
        args.out,
        name=args.name,
        motions=motions,
        skip_gate=args.skip_gate,
    )
    print(f"done in {time.time() - t0:.0f}s", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())