import os, shutil, subprocess, tempfile

SUPPORTED = (".glb", ".gltf", ".fbx", ".dae")
_SCRIPT = os.path.join(os.path.dirname(__file__), "tools", "fbx2glb.mjs")
_NODE_DIR = os.environ.get("VID2RIG_NODE_DIR", "/tmp/vid2rig_node")

def _find(cmd):
    """Locate a binary even if it's not on the Python process PATH (HF runtime PATH != build PATH)."""
    return shutil.which(cmd) or next(
        (p for p in (f"/usr/local/bin/{cmd}", f"/usr/bin/{cmd}", f"/opt/homebrew/bin/{cmd}")
         if os.path.exists(p)), None)

def _node_convert(src: str, out: str) -> bool:
    """FBX -> GLB via three.js FBXLoader + GLTFExporter (the path the mule scenes use). Correctly
    collapses Mixamo $AssimpFbx$ pivots, fixes the up-axis, and writes TRS bones so the baked
    animation is valid. Installs `three` into a work dir on first use. Returns True on success."""
    node, npm = _find("node"), _find("npm")
    if not node or not npm or not os.path.exists(_SCRIPT):
        print(f"[rig_io] node convert unavailable (node={node} npm={npm} script={os.path.exists(_SCRIPT)})", flush=True)
        return False
    src, out = os.path.abspath(src), os.path.abspath(out)   # node runs with cwd=_NODE_DIR
    os.makedirs(_NODE_DIR, exist_ok=True)
    shutil.copy(_SCRIPT, os.path.join(_NODE_DIR, "fbx2glb.mjs"))
    if not os.path.isdir(os.path.join(_NODE_DIR, "node_modules", "three")):
        r = subprocess.run([npm, "install", "--no-audit", "--no-fund", "three@0.169.0"],
                           cwd=_NODE_DIR, capture_output=True, text=True)
        if r.returncode != 0:
            print("[rig_io] npm install three FAILED:", (r.stderr or r.stdout)[-500:], flush=True)
            return False
    r = subprocess.run([node, os.path.join(_NODE_DIR, "fbx2glb.mjs"), src, out],
                       cwd=_NODE_DIR, capture_output=True, text=True)
    ok = r.returncode == 0 and os.path.exists(out)
    if ok:
        print("[rig_io] three.js conversion OK ->", out, flush=True)
    else:
        print("[rig_io] three.js conversion FAILED:", (r.stderr or r.stdout)[-500:], flush=True)
    return ok

def _embedded_png(fbx: str):
    """(bytes, mime) of the FBX's first embedded texture via assimp, or None."""
    assimp = _find("assimp")
    if not assimp:
        return None
    try:
        import pygltflib as gl
        tmp = tempfile.mktemp(suffix=".glb")
        subprocess.run([assimp, "export", fbx, tmp], check=True, capture_output=True)
        ga = gl.GLTF2().load(tmp)
        if not ga.images:
            return None
        im = ga.images[0]; bv = ga.bufferViews[im.bufferView]
        png = (ga.binary_blob() or b"")[(bv.byteOffset or 0):(bv.byteOffset or 0) + bv.byteLength]
        return png, (im.mimeType or "image/png")
    except Exception:
        return None

def _finalize_rig(glb: str, fbx: str) -> None:
    """Post-process the clean three.js GLB in place: (1) point skins[0].skeleton at Hips
    (fixes SKIN_SKELETON_INVALID), (2) transplant the FBX's embedded texture onto the material
    (three.js export drops it but keeps the matching UV0). Best-effort; failures leave a valid rig."""
    try:
        import pygltflib as gl
        g = gl.GLTF2().load(glb)
        changed = False
        if g.skins:
            hips = next((i for i, n in enumerate(g.nodes)
                         if n.name and n.name.replace(":", "").replace("_", "") == "mixamorigHips"), None)
            if hips is not None and g.skins[0].skeleton != hips:
                g.skins[0].skeleton = hips; changed = True

        tex = _embedded_png(fbx)
        prims = [p for m in (g.meshes or []) for p in m.primitives
                 if getattr(p.attributes, "TEXCOORD_0", None) is not None]
        if tex and prims:
            png, mime = tex
            blob = bytearray(g.binary_blob() or b"")
            while len(blob) % 4:
                blob.append(0)
            off = len(blob); blob += png
            g.bufferViews.append(gl.BufferView(buffer=0, byteOffset=off, byteLength=len(png)))
            g.images = (g.images or []) + [gl.Image(bufferView=len(g.bufferViews) - 1, mimeType=mime)]
            g.samplers = (g.samplers or []) + [gl.Sampler()]
            g.textures = (g.textures or []) + [gl.Texture(source=len(g.images) - 1, sampler=len(g.samplers) - 1)]
            tex_i = len(g.textures) - 1
            for prim in prims:
                if prim.material is None:
                    g.materials = (g.materials or []) + [gl.Material()]; prim.material = len(g.materials) - 1
                mat = g.materials[prim.material]
                if mat.pbrMetallicRoughness is None:
                    mat.pbrMetallicRoughness = gl.PbrMetallicRoughness()
                mat.pbrMetallicRoughness.baseColorTexture = gl.TextureInfo(index=tex_i, texCoord=0)
                mat.pbrMetallicRoughness.baseColorFactor = [1, 1, 1, 1]
            g.set_binary_blob(bytes(blob)); g.buffers[0].byteLength = len(blob)
            changed = True
            print("[rig_io] texture transplanted onto rig", flush=True)
        if changed:
            g.save(glb)
    except Exception as e:
        print("[rig_io] _finalize_rig error:", repr(e), flush=True)
        return

def to_glb(path: str) -> str:
    """Return a .glb path for the rig. .glb/.gltf pass through; .fbx is converted via three.js
    (preferred — handles Mixamo rigs correctly); assimp is a last-resort fallback (it mangles
    Mixamo pivots, so only used when node/three is unavailable, and for .dae)."""
    ext = os.path.splitext(path)[1].lower()
    if ext in (".glb", ".gltf"):
        return path
    if ext not in SUPPORTED:
        raise ValueError(f"unsupported rig format {ext!r}; use one of {SUPPORTED}")
    out = tempfile.mktemp(suffix=".glb")
    if ext == ".fbx" and _node_convert(path, out):
        _finalize_rig(out, path)          # fix skin root + borrow embedded texture
        return out
    print("[rig_io] FALLING BACK to assimp (mangles Mixamo rig + no texture)", flush=True)
    assimp = _find("assimp")
    if assimp:
        subprocess.run([assimp, "export", path, out], check=True, capture_output=True)
        return out
    fbx2gltf = shutil.which("FBX2glTF") or shutil.which("fbx2gltf")
    if fbx2gltf and ext == ".fbx":
        subprocess.run([fbx2gltf, "-i", path, "-o", out, "--binary"], check=True, capture_output=True)
        return out if os.path.exists(out) else out + ".glb"
    raise RuntimeError(
        "No FBX/DAE converter available. Need Node.js + three (preferred) or `brew install assimp`.")