"""Minimal public ZeroGPU shell for the private BeyondRetarget-plus runtime."""

from __future__ import annotations

import os
import shutil
import subprocess
import sys
from pathlib import Path

import gradio as gr
import spaces
from huggingface_hub import snapshot_download


PRIVATE_REPO_ID = os.environ.get(
    "PRIVATE_REPO_ID",
    "bear-ty/BeyondRetargetPlus-runtime-private",
).strip()
PRIVATE_REPO_TOKEN = os.environ.get("PRIVATE_REPO_TOKEN", "").strip()

if not PRIVATE_REPO_TOKEN:
    raise RuntimeError(
        "PRIVATE_REPO_TOKEN is not configured. Add it as a Hugging Face Space secret."
    )

PKG_ROOT = Path("/tmp/brplus_pkg")
RUNTIME_ROOT = Path("/tmp/runtime")
if not (RUNTIME_ROOT / "patch_runtime.py").is_file() or os.environ.get("BRPLUS_FORCE_REFRESH", "1") == "1":
    snapshot_download(
        repo_id=PRIVATE_REPO_ID,
        repo_type="model",
        token=PRIVATE_REPO_TOKEN,
        local_dir=str(PKG_ROOT),
        allow_patterns=["pkg/*", "overlay/*"],
    )
    parts = sorted((PKG_ROOT / "pkg").glob("runtime_part_*"))
    if not parts:
        raise RuntimeError("Private Plus runtime package is missing.")
    command = "cat " + " ".join(str(part) for part in parts) + " | tar -xf - -C /tmp"
    subprocess.run(command, shell=True, check=True)
    overlay = PKG_ROOT / "overlay"
    if overlay.is_dir():
        for item in overlay.iterdir():
            if item.is_file():
                shutil.copy2(item, RUNTIME_ROOT / item.name)
    shutil.rmtree(PKG_ROOT, ignore_errors=True)

RUNTIME_ROOT = RUNTIME_ROOT.resolve()
patch_script = RUNTIME_ROOT / "patch_runtime.py"
if patch_script.is_file():
    subprocess.run([sys.executable, str(patch_script)], check=True)

required = [
    RUNTIME_ROOT / "app_runtime.py",
    RUNTIME_ROOT / "plus_engine.py",
    RUNTIME_ROOT / "weights" / "vimo_checkpoint.pth.tar",
    RUNTIME_ROOT / "weights" / "experiments" /
    "stage2_shared_static" / "20261017_012204" / "checkpoints" / "best.pt",
    RUNTIME_ROOT / "weights" / "yolo" / "yolov8x.pt",
    RUNTIME_ROOT / "vitpose" / "inputs" / "checkpoints" / "vitpose" / "vitpose-h-multi-coco.pth",
    RUNTIME_ROOT / "src" / "configs" / "config_beyondretarget_motion_g1.yaml",
    RUNTIME_ROOT / "renderer" / "render_8robot_two_rows_remote.py",
]
missing = [str(path.relative_to(RUNTIME_ROOT)) for path in required if not path.is_file()]
if missing:
    raise RuntimeError(f"Private Plus runtime is incomplete; missing: {', '.join(missing)}")

os.environ.setdefault("BRPLUS_SRC", str(RUNTIME_ROOT / "src"))
os.environ.setdefault("BRPLUS_WEIGHTS", str(RUNTIME_ROOT / "weights"))
os.environ.setdefault("BRPLUS_RENDERER", str(RUNTIME_ROOT / "renderer" / "render_8robot_two_rows_remote.py"))
os.environ.setdefault("BRPLUS_RENDER_XML", str(RUNTIME_ROOT / "renderer" / "assets" / "atlas_g1_gr1t1_h1_gr2_r1_tk_t1.xml"))
os.environ.setdefault("BRPLUS_VITPOSE_ROOT", str(RUNTIME_ROOT / "vitpose"))
os.environ.setdefault("BRPLUS_OUTPUT_DIR", "/tmp/brplus_outputs")
os.environ.setdefault("MUJOCO_GL", "egl")
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
os.environ.setdefault("LIBGL_ALWAYS_SOFTWARE", "1")
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")

sys.path.insert(0, str(RUNTIME_ROOT))

from app_runtime import create_demo, run_gpu_stage  # noqa: E402
import cv2  # noqa: E402


@spaces.GPU(duration=80, size="large")
def gpu_fast(video_path: str, mode: str, progress=gr.Progress()) -> dict:
    return run_gpu_stage(video_path, mode, progress=progress)


@spaces.GPU(duration=180, size="large")
def gpu_full(video_path: str, mode: str, progress=gr.Progress()) -> dict:
    return run_gpu_stage(video_path, mode, progress=progress)


def zero_gpu_inference(video_path: str, mode: str, progress=gr.Progress()) -> dict:
    capture = cv2.VideoCapture(video_path)
    try:
        fps = float(capture.get(cv2.CAP_PROP_FPS) or 0.0)
        frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    finally:
        capture.release()
    duration = frames / fps if fps > 0 else 0.0
    # Short clips fit the anonymous 120s ZeroGPU allowance; longer clips use
    # the 225s allowance, which requires a signed-in Hugging Face user.
    if duration <= 2.2:
        return gpu_fast(video_path, mode, progress=progress)
    return gpu_full(video_path, mode, progress=progress)


demo = create_demo(zero_gpu_inference)
demo.queue(default_concurrency_limit=1, max_size=3)


if __name__ == "__main__":
    demo.launch()