"""SAM 3D Body runner (GPU only). Clones facebookresearch/sam-3d-body + detectron2 at first use,
loads the estimator, runs per-frame inference. Heavy imports (torch/cv2/the repo) stay in here so
importing vid2rig never needs them. Designed for a persistent-GPU Space (A10G); needs HF_TOKEN with
gated access to facebook/sam-3d-body-dinov3."""
import os, sys, subprocess, functools
import numpy as np

SAM3D_DIR = os.environ.get("SAM3D_BODY_ROOT", "/tmp/sam-3d-body")
DETECTRON2 = "git+https://github.com/facebookresearch/detectron2.git@a1ce2f9"

def _pip(*args):
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", *args], check=True)

def _ensure_repo():
    if not os.path.isdir(SAM3D_DIR):
        subprocess.run(["git", "clone", "--depth", "1",
                        "https://github.com/facebookresearch/sam-3d-body.git", SAM3D_DIR], check=True)
    if SAM3D_DIR not in sys.path:
        sys.path.insert(0, SAM3D_DIR)
    try:
        import detectron2  # noqa: F401
    except Exception:
        _pip(DETECTRON2, "--no-build-isolation", "--no-deps")

@functools.lru_cache(maxsize=1)
def _estimator():
    _ensure_repo()
    from notebook.utils import setup_sam_3d_body
    return setup_sam_3d_body(hf_repo_id="facebook/sam-3d-body-dinov3")

def _log_structure(outs):
    print("[sam3d] people detected:", len(outs), flush=True)
    if not outs:
        return
    o = outs[0]
    print("[sam3d] output keys:", list(o.keys()), flush=True)
    for k, v in o.items():
        shp = getattr(v, "shape", None)
        print(f"[sam3d]   {k}: shape={tuple(shp) if shp is not None else type(v).__name__}", flush=True)

def _to_list(x):
    try:
        return x.detach().cpu().numpy().tolist()
    except Exception:
        try:
            return list(x)
        except Exception:
            return x

def _get_skeleton(est):
    """Return the MHR skeleton {names, parents, prerotations, offsets} from the loaded model."""
    out = {"names": None, "parents": None, "prerotations": None, "offsets": None}
    try:
        char = None
        for path in ("model.head_pose.mhr.character_torch", "model.head_pose_hand.mhr.character_torch",
                     "head_pose.mhr.character_torch"):
            cur = est
            for p in path.split("."):
                cur = getattr(cur, p, None)
                if cur is None:
                    break
            if cur is not None:
                char = cur; break
        if char is None:
            print("[sam3d-skel] character not found", flush=True); return out
        skel = getattr(char, "skeleton", char)
        names = getattr(skel, "joint_names", None) or getattr(char, "joint_names", None)
        if names is not None:
            out["names"] = [str(n) for n in names]
        jp = getattr(skel, "joint_parents", None)
        if jp is not None:
            out["parents"] = [int(x) for x in _to_list(jp)]
        pr = getattr(skel, "joint_prerotations", None)
        if pr is not None:
            out["prerotations"] = np.asarray(_to_list(pr), float)
            print("[sam3d-skel] prerotations shape:", out["prerotations"].shape,
                  "row0:", out["prerotations"].reshape(out["prerotations"].shape[0], -1)[0].tolist(), flush=True)
        off = getattr(skel, "joint_translation_offsets", None)
        if off is not None:
            out["offsets"] = np.asarray(_to_list(off), float)
        print(f"[sam3d-skel] names={None if out['names'] is None else len(out['names'])} "
              f"parents={None if out['parents'] is None else len(out['parents'])}", flush=True)
    except Exception as e:
        print("[sam3d-skel] introspection failed:", repr(e), flush=True)
    return out

def run_sam3d_frames(video_path, max_frames=120, log=True):
    """Return (global_rots [F,127,3,3], skeleton dict, src_fps). Uses the primary person per frame."""
    import cv2
    _ensure_repo()
    est = _estimator()
    skel = _get_skeleton(est)
    faces = None
    try:
        faces = np.asarray(_to_list(getattr(est, "faces", None)))
    except Exception:
        pass
    cap = cv2.VideoCapture(video_path)
    src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    rots, mesh, i = [], None, 0
    while i < max_frames:
        ok, bgr = cap.read()
        if not ok:
            break
        outs = est.process_one_image(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
        if log and i == 0 and outs:
            _log_structure(outs)
        if outs and outs[0].get("pred_global_rots") is not None:
            rots.append(np.asarray(_to_list(outs[0]["pred_global_rots"])))
            if mesh is None and faces is not None and outs[0].get("pred_vertices") is not None:
                mesh = (np.asarray(_to_list(outs[0]["pred_vertices"])), faces)   # frame-0 showcase mesh
        i += 1
    cap.release()
    g = np.stack(rots) if rots else np.zeros((0, 0, 3, 3))
    return g, skel, src_fps, mesh