#!/usr/bin/env python3
"""
hyworld_install.py — Install HY-World-2.0 into an isolated venv.

This script:
  1. Performs preflight checks (disk space, GPU VRAM)
  2. Clones / updates the HY-World-2.0 GitHub repo
  3. Creates an isolated venv (Python 3.11+)
  4. Installs base dependencies (requirements.txt only, not the heavy world-gen pipeline)
  5. Best-effort: installs gsplat_maskgaussian and flash-attn (fails gracefully if not available)
  6. Writes a capabilities manifest and final result sentinel to stdout

All output is stdlib only — no third-party deps, so this can run with any ambient Python.
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import pathlib
import platform
import re
import shutil
import subprocess
import sys
import tempfile
from typing import Optional

log = logging.getLogger("hyworld_install")


def _get_venv_python(venv_path: pathlib.Path) -> pathlib.Path:
    """Return the python executable path for a venv."""
    if platform.system() == "Windows":
        return venv_path / ".venv" / "Scripts" / "python.exe"
    return venv_path / ".venv" / "bin" / "python"


def _pip_env(install_path: pathlib.Path) -> dict:
    """
    Build an environment for pip subprocess calls that redirects TEMP/TMP and
    pip's own cache to a directory on the SAME drive as the install path.
    Model/package downloads here can be many GB, and pip's default temp/cache
    dirs live under the system TEMP (usually the C: drive) regardless of where
    the venv itself lives — a small system drive fills up and aborts the
    install even when the target drive has plenty of room.
    """
    tmp_dir = install_path / ".pip_tmp"
    cache_dir = install_path / ".pip_cache"
    tmp_dir.mkdir(parents=True, exist_ok=True)
    cache_dir.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    env["TEMP"] = str(tmp_dir)
    env["TMP"] = str(tmp_dir)
    env["TMPDIR"] = str(tmp_dir)
    env["PIP_CACHE_DIR"] = str(cache_dir)
    return env


def _existing_ancestor(path: pathlib.Path) -> pathlib.Path:
    """Walk up to the nearest existing directory (path itself may not exist yet)."""
    p = path
    while not p.exists() and p.parent != p:
        p = p.parent
    return p


def _check_disk_space(path: pathlib.Path, required_gb: int = 20) -> bool:
    """Check if the target path AND the system temp drive have enough free space.
    Package/model downloads stage through the system temp dir (often a different,
    smaller drive than the install path), so both need checking. Warns, doesn't block."""
    ok = True
    try:
        stat = shutil.disk_usage(_existing_ancestor(path))
        free_gb = stat.free / (1024 ** 3)
        if free_gb < required_gb:
            log.warning(
                "Only %.1f GB free at install path; HY-World-2.0 needs ~%d GB. Proceeding anyway.",
                free_gb, required_gb,
            )
            ok = False
        else:
            log.info("Install path disk space OK: %.1f GB free", free_gb)
    except Exception as e:
        log.warning("Could not check install path disk space: %s", e)

    try:
        temp_stat = shutil.disk_usage(_existing_ancestor(pathlib.Path(tempfile.gettempdir())))
        temp_free_gb = temp_stat.free / (1024 ** 3)
        if temp_free_gb < 5:
            log.warning(
                "System temp drive (%s) has only %.1f GB free. Downloads are redirected to the "
                "install path's own temp dir to avoid filling this drive, but other tools sharing "
                "it may still be affected.",
                tempfile.gettempdir(), temp_free_gb,
            )
            ok = False
        else:
            log.info("System temp drive OK: %.1f GB free", temp_free_gb)
    except Exception as e:
        log.warning("Could not check system temp drive space: %s", e)

    return ok


def _check_gpu_vram() -> Optional[int]:
    """Check available GPU VRAM in GB using nvidia-smi. Returns GB or None if unavailable."""
    try:
        result = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=memory.total",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0:
            mb = int(result.stdout.strip().split()[0])
            gb = mb / 1024
            if gb < 8:
                log.warning(
                    "GPU VRAM: %.1f GB detected. HY-Pano-2.0 (text/image stage) is unlikely to fit; "
                    "WorldMirror-2.0 (3D reconstruction) should work.",
                    gb,
                )
            else:
                log.info("GPU VRAM: %.1f GB available", gb)
            return gb
    except (FileNotFoundError, ValueError, subprocess.TimeoutExpired):
        pass
    log.info("GPU VRAM: could not detect (nvidia-smi not found or timed out)")
    return None


def _clone_or_pull(repo_url: str, target_path: pathlib.Path) -> bool:
    """Clone or pull the HY-World-2.0 repo. Returns True on success."""
    should_clone = False
    if target_path.exists():
        log.info("Repo directory exists at %s, checking if valid git repo…", target_path)
        try:
            subprocess.run(
                ["git", "-C", str(target_path), "pull"],
                check=True,
                capture_output=True,
            )
            log.info("Pulled latest changes")
            return True
        except subprocess.CalledProcessError as e:
            log.warning("Not a valid git repo or pull failed: %s. Removing and cloning fresh…", e)
            try:
                shutil.rmtree(target_path)
                should_clone = True
            except Exception as e2:
                log.error("Failed to remove directory: %s", e2)
                return False
    else:
        should_clone = True

    if should_clone:
        log.info("Cloning %s → %s…", repo_url, target_path)
        try:
            subprocess.run(
                ["git", "clone", repo_url, str(target_path)],
                check=True,
                capture_output=False,
            )
            return True
        except subprocess.CalledProcessError as e:
            log.error("Failed to clone: %s", e)
            return False

    return False


def _create_venv(venv_path: pathlib.Path) -> bool:
    """Create an isolated Python venv. Returns True on success."""
    venv_dir = venv_path / ".venv"
    if venv_dir.exists():
        log.info("Venv already exists at %s", venv_dir)
        return True

    log.info("Creating venv at %s…", venv_dir)
    try:
        subprocess.run(
            [sys.executable, "-m", "venv", str(venv_dir)],
            check=True,
            capture_output=False,
        )
        return True
    except subprocess.CalledProcessError as e:
        log.error("Failed to create venv: %s", e)
        return False


def _pip_install(
    venv_python: pathlib.Path, cwd: pathlib.Path, requirements_file: str, env: dict
) -> bool:
    """Install from a requirements file. Returns True on success."""
    req_path = cwd / requirements_file
    if not req_path.exists():
        log.warning("Requirements file not found: %s", req_path)
        return False

    log.info("Installing %s…", requirements_file)
    try:
        subprocess.run(
            [str(venv_python), "-m", "pip", "install", "-r", str(req_path), "--no-build-isolation"],
            cwd=str(cwd),
            check=True,
            capture_output=False,
            env=env,
        )
        return True
    except subprocess.CalledProcessError as e:
        # If install failed, try without CuPy (which often fails on Windows without CUDA SDK)
        log.warning("Initial install failed; trying without CuPy…")
        req_text = req_path.read_text(encoding="utf-8")
        cupy_line = [l for l in req_text.split('\n') if 'cupy' in l.lower()]
        if cupy_line:
            log.warning("Skipping CuPy: %s (needs CUDA SDK; will use CPU fallbacks)", cupy_line[0])
            # Write a temp requirements file without CuPy
            temp_req = req_path.parent / f"{req_path.stem}_nocupy{req_path.suffix}"
            filtered = '\n'.join(l for l in req_text.split('\n') if 'cupy' not in l.lower())
            temp_req.write_text(filtered, encoding="utf-8")
            try:
                subprocess.run(
                    [str(venv_python), "-m", "pip", "install", "-r", str(temp_req), "--no-build-isolation"],
                    cwd=str(cwd),
                    check=True,
                    capture_output=False,
                    env=env,
                )
                temp_req.unlink()
                return True
            except subprocess.CalledProcessError as e2:
                log.error("pip install (no-CuPy) also failed: %s", e2)
                return False
        else:
            log.error("pip install failed: %s", e)
            return False


def _pip_install_editable(
    venv_python: pathlib.Path, package_path: pathlib.Path, env: dict
) -> bool:
    """Install a package in editable mode. Returns True on success."""
    if not package_path.exists():
        log.warning("Package path does not exist: %s", package_path)
        return False

    log.info("Installing editable: %s…", package_path)
    try:
        subprocess.run(
            [
                str(venv_python),
                "-m",
                "pip",
                "install",
                "-e",
                str(package_path),
                "--no-build-isolation",
            ],
            check=True,
            capture_output=False,
            env=env,
        )
        return True
    except subprocess.CalledProcessError as e:
        log.error("Editable install failed: %s", e)
        return False


def _pip_install_package(
    venv_python: pathlib.Path, package_name: str, env: dict
) -> bool:
    """Install a package by name. Returns True on success."""
    log.info("Installing %s…", package_name)
    try:
        subprocess.run(
            [
                str(venv_python),
                "-m",
                "pip",
                "install",
                package_name,
                "--no-build-isolation",
            ],
            check=True,
            capture_output=False,
            env=env,
        )
        return True
    except subprocess.CalledProcessError as e:
        log.error("Package install failed: %s", e)
        return False


# HY-World-2.0's worldrecon attention module does:
#   try: from flash_attn_interface import flash_attn_func as flash_attn_func_v3
#   except ImportError: from flash_attn.flash_attn_interface import flash_attn_func as flash_attn_func_v2
# ...with NO further fallback — if flash_attn isn't installed, the whole module
# fails to import, even though the function is only actually *called* for
# fp16/bf16 tensors (the fp32/CPU path already uses F.scaled_dot_product_attention
# instead, per hyworld2/worldrecon/hyworldmirror/models/layers/attention.py).
# On a machine without a CUDA toolchain, the real flash-attn package can't be
# built. This stub satisfies the import without needing compiled CUDA kernels;
# it only raises if the accelerated path is actually reached (which won't
# happen when running fp32/CPU-only, as this installer targets).
_FLASH_ATTN_STUB = '''"""Stub flash_attn — satisfies imports on machines without a CUDA toolchain.
Only raises if the accelerated fp16/bf16 code path is actually reached."""


def flash_attn_func(*args, **kwargs):
    raise NotImplementedError(
        "flash_attn is not installed (no CUDA toolchain available). "
        "This code path requires fp16/bf16 tensors on a real GPU with flash-attn built; "
        "run in fp32 (CPU-compatible) mode instead."
    )
'''


def _venv_site_packages(venv_python: pathlib.Path, env: dict) -> pathlib.Path:
    """Resolve the venv's actual site-packages dir (index [0] is the venv root
    on some platforms, not site-packages — filter for the entry containing it)."""
    code = "import site; print([p for p in site.getsitepackages() if 'site-packages' in p][0])"
    out = subprocess.run(
        [str(venv_python), "-c", code], check=True, capture_output=True, text=True, env=env,
    ).stdout.strip()
    return pathlib.Path(out)


def _install_flash_attn_stub(venv_python: pathlib.Path, env: dict) -> bool:
    """Install a minimal stub flash_attn package into the venv's site-packages."""
    try:
        pkg_dir = _venv_site_packages(venv_python, env) / "flash_attn"
        pkg_dir.mkdir(parents=True, exist_ok=True)
        (pkg_dir / "__init__.py").write_text(_FLASH_ATTN_STUB, encoding="utf-8")
        (pkg_dir / "flash_attn_interface.py").write_text(_FLASH_ATTN_STUB, encoding="utf-8")
        log.info("Installed flash_attn stub at %s (import-compatible, CPU/fp32 only)", pkg_dir)
        return True
    except Exception as e:
        log.warning("Could not install flash_attn stub: %s", e)
        return False


# hyworld2/worldrecon/hyworldmirror/models/models/worldmirror.py unconditionally
# does `from .rasterization import GaussianSplatRenderer`, and rasterization.py
# unconditionally does `from gsplat.rendering import rasterization` +
# `from gsplat.strategy import DefaultStrategy` — so importing the WorldMirror
# model class at all requires gsplat, even when the "gs" (Gaussian Splat) output
# head is disabled via disable_heads and no gsplat function is ever called.
# gsplat's real package needs compiled CUDA kernels (same toolchain problem as
# flash-attn); this stub satisfies the import only.
_GSPLAT_RENDERING_STUB = '''"""Stub gsplat.rendering — satisfies imports when the GS head is disabled."""


def rasterization(*args, **kwargs):
    raise NotImplementedError(
        "gsplat is not installed (no CUDA toolchain available). "
        "Gaussian Splat output is unavailable; disable the 'gs' head to avoid this path."
    )
'''

_GSPLAT_STRATEGY_STUB = '''"""Stub gsplat.strategy — satisfies imports when the GS head is disabled."""


class DefaultStrategy:
    def __init__(self, *args, **kwargs):
        pass
'''


def _install_gsplat_stub(venv_python: pathlib.Path, env: dict) -> bool:
    """Install a minimal stub gsplat package into the venv's site-packages."""
    try:
        pkg_dir = _venv_site_packages(venv_python, env) / "gsplat"
        pkg_dir.mkdir(parents=True, exist_ok=True)
        (pkg_dir / "__init__.py").write_text("", encoding="utf-8")
        (pkg_dir / "rendering.py").write_text(_GSPLAT_RENDERING_STUB, encoding="utf-8")
        (pkg_dir / "strategy.py").write_text(_GSPLAT_STRATEGY_STUB, encoding="utf-8")
        log.info("Installed gsplat stub at %s (import-compatible; GS output disabled)", pkg_dir)
        return True
    except Exception as e:
        log.warning("Could not install gsplat stub: %s", e)
        return False


def _enable_gpu(venv_python: pathlib.Path, pip_env: dict, capabilities: dict, cuda: str) -> None:
    """Upgrade an existing hyworld venv for real GPU use, in-place:

    1. Swap torch (and torchvision, if present) for the matching +cuXXX builds —
       pinned to the exact versions already in the venv, because a bare
       `torch==X.Y.Z` is a no-op (pip treats the +cpu build as satisfying it).
    2. Try a REAL gsplat from its prebuilt-wheel index (docs.gsplat.studio/whl)
       — the only viable route on machines without nvcc/MSVC, where the earlier
       full install correctly fell back to an import-stub. The stale stub dir
       (written directly into site-packages, no dist-info) is removed first so
       a real wheel doesn't overlay it.
    3. flash-attn: only attempted when nvcc exists — there are no official
       Windows wheels and a source build without a CUDA toolchain cannot work,
       so without nvcc the honest answer is to keep the stub and say so.

    Mutates `capabilities` (torch_cuda / gsplat / flash_attn) to reflect what
    ACTUALLY happened — never records success for a fallback.
    """
    def probe(pkg: str):
        r = subprocess.run(
            [str(venv_python), "-m", "pip", "show", pkg],
            capture_output=True, text=True, env=pip_env,
        )
        m = re.search(r"^Version:\s*(\d+\.\d+\.\d+)", r.stdout, re.M)
        return m.group(1) if m else None

    log.info("─── GPU: CUDA torch ───")
    torch_ver = probe("torch")
    if not torch_ver:
        log.warning("torch is not installed in the hyworld venv — run a full install first")
        capabilities["torch_cuda"] = False
        return
    specs = [f"torch=={torch_ver}+{cuda}"]
    tv_ver = probe("torchvision")
    if tv_ver:
        specs.append(f"torchvision=={tv_ver}+{cuda}")
    log.info("installing %s (multi-GB download)…", " ".join(specs))
    try:
        subprocess.run(
            [str(venv_python), "-m", "pip", "install", "--timeout", "120", *specs,
             "--index-url", f"https://download.pytorch.org/whl/{cuda}"],
            check=True, env=pip_env,
        )
    except subprocess.CalledProcessError:
        log.warning("✗ CUDA torch install failed (no %s build for torch %s? try --cuda cu128)", cuda, torch_ver)
        capabilities["torch_cuda"] = False
        return
    r = subprocess.run(
        [str(venv_python), "-c", "import torch;print(torch.cuda.is_available())"],
        capture_output=True, text=True, env=pip_env,
    )
    cuda_ok = r.stdout.strip().endswith("True")
    capabilities["torch_cuda"] = cuda if cuda_ok else False
    log.info("%s torch CUDA: %s", "✓" if cuda_ok else "✗", r.stdout.strip())

    log.info("─── GPU: real gsplat (prebuilt wheel) ───")
    if capabilities.get("gsplat") == "stub":
        try:
            stub_dir = _venv_site_packages(venv_python, pip_env) / "gsplat"
            if stub_dir.exists() and not (stub_dir.parent / "gsplat-0.0.0.dist-info").exists():
                shutil.rmtree(stub_dir)
                log.info("removed gsplat import-stub before real install")
        except Exception as e:
            log.warning("could not remove gsplat stub: %s", e)
    # gsplat's wheel index is per torch-minor + CUDA tag, e.g. pt27cu126.
    tag = f"pt{''.join(torch_ver.split('.')[:2])}{cuda}"
    real_gsplat = False
    try:
        subprocess.run(
            [str(venv_python), "-m", "pip", "install", "--timeout", "120", "gsplat",
             "--index-url", f"https://docs.gsplat.studio/whl/{tag}"],
            check=True, env=pip_env,
        )
        r = subprocess.run(
            [str(venv_python), "-c", "from gsplat.rendering import rasterization; import gsplat; print(getattr(gsplat,'__version__','?'))"],
            capture_output=True, text=True, env=pip_env,
        )
        real_gsplat = r.returncode == 0
    except subprocess.CalledProcessError:
        pass
    if real_gsplat:
        capabilities["gsplat"] = True
        log.info("✓ real gsplat installed (Gaussian Splat output head available)")
    else:
        log.warning(
            "✗ no prebuilt gsplat wheel for %s and no local CUDA toolchain to build one — restoring import-stub. "
            "(gsplat's wheel index currently covers torch 2.0-2.4 + cu118/cu121/cu124 on Python 3.10 only; "
            "outside that matrix, real gsplat needs the CUDA toolkit installed so it can compile.)", tag)
        if _install_gsplat_stub(venv_python, pip_env):
            capabilities["gsplat"] = "stub"

    log.info("─── GPU: flash-attn ───")
    if shutil.which("nvcc"):
        if _pip_install_package(venv_python, "flash-attn", pip_env):
            capabilities["flash_attn"] = True
            log.info("✓ flash-attn built against local CUDA toolchain")
        else:
            log.warning("✗ flash-attn build failed — keeping stub")
    else:
        log.info("no nvcc on PATH — flash-attn has no official Windows wheels and cannot be "
                 "built without the CUDA toolkit; keeping the import-stub (fp32 paths still work, "
                 "now GPU-accelerated via CUDA torch)")


def main():
    ap = argparse.ArgumentParser(
        description="Install HY-World-2.0 into an isolated venv"
    )
    ap.add_argument(
        "--path",
        default=str(pathlib.Path.home() / ".zelpi" / "hyworld2"),
        help="Installation path (default: ~/.zelpi/hyworld2)",
    )
    ap.add_argument(
        "--skip-gsplat", action="store_true", help="Skip gsplat_maskgaussian build"
    )
    ap.add_argument(
        "--skip-flash-attn", action="store_true", help="Skip flash-attn install"
    )
    ap.add_argument(
        "--skip-pybullet", action="store_true", help="Skip pybullet install"
    )
    ap.add_argument(
        "--pybullet-only", action="store_true",
        help="Only install pybullet into an existing venv (skip clone/venv/base deps)",
    )
    ap.add_argument(
        "--gpu-only", action="store_true",
        help="Upgrade an existing venv for GPU: CUDA torch + real gsplat wheel + flash-attn where possible",
    )
    ap.add_argument(
        "--cuda", default="cu126",
        help="CUDA wheel channel for --gpu-only (default: cu126)",
    )
    args = ap.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s  %(name)s  %(levelname)s  %(message)s",
        datefmt="%H:%M:%S",
    )

    install_path = pathlib.Path(args.path).expanduser()
    repo_url = "https://github.com/Tencent-Hunyuan/HY-World-2.0"

    # Load any existing capabilities manifest so a --pybullet-only (or any
    # partial) run merges into it instead of clobbering previously-recorded
    # gsplat/flash_attn/pybullet status.
    manifest_path = install_path / ".zelpi_capabilities.json"
    capabilities = {"gsplat": False, "flash_attn": False, "pybullet": False, "python": None, "path": str(install_path)}
    if manifest_path.exists():
        try:
            capabilities.update(json.loads(manifest_path.read_text(encoding="utf-8")))
        except Exception:
            pass
    result = {"code": "success", "capabilities": capabilities}

    if args.gpu_only:
        venv_python = _get_venv_python(install_path)
        if not venv_python.exists():
            result["code"] = "venv_missing"
            result["message"] = f"No existing venv at {install_path} — run a full install first."
            print("HYWORLD_INSTALL_RESULT " + json.dumps(result))
            sys.exit(1)
        capabilities["python"] = str(venv_python)
        pip_env = _pip_env(install_path)
        if _check_gpu_vram() is None:
            result["code"] = "no_gpu"
            result["message"] = "no NVIDIA GPU detected (nvidia-smi missing/unresponsive)"
            print("HYWORLD_INSTALL_RESULT " + json.dumps(result))
            sys.exit(1)
        _enable_gpu(venv_python, pip_env, capabilities, args.cuda)
        manifest_path.write_text(json.dumps(capabilities, indent=2), encoding="utf-8")
        result["capabilities"] = capabilities
        print("HYWORLD_INSTALL_RESULT " + json.dumps(result))
        return

    if args.pybullet_only:
        venv_python = _get_venv_python(install_path)
        if not venv_python.exists():
            result["code"] = "venv_missing"
            result["message"] = f"No existing venv at {install_path} — run a full install first."
            print("HYWORLD_INSTALL_RESULT " + json.dumps(result))
            sys.exit(1)
        capabilities["python"] = str(venv_python)
        pip_env = _pip_env(install_path)
        log.info("─── PYBULLET ONLY ───")
        if _pip_install_package(venv_python, "pybullet", pip_env):
            capabilities["pybullet"] = True
            log.info("✓ pybullet installed")
        else:
            capabilities["pybullet"] = False
            log.warning("✗ pybullet install failed")
        manifest_path.write_text(json.dumps(capabilities, indent=2), encoding="utf-8")
        result["capabilities"] = capabilities
        print("HYWORLD_INSTALL_RESULT " + json.dumps(result))
        return

    # Preflight checks
    log.info("─── PREFLIGHT ───")
    _check_disk_space(install_path, required_gb=20)
    _check_gpu_vram()

    # Clone or pull repo
    log.info("─── CLONE/PULL REPO ───")
    if not _clone_or_pull(repo_url, install_path):
        result["code"] = "clone_failed"
        print("HYWORLD_INSTALL_RESULT " + json.dumps(result))
        sys.exit(1)

    # Create venv
    log.info("─── CREATE VENV ───")
    if not _create_venv(install_path):
        result["code"] = "venv_failed"
        print("HYWORLD_INSTALL_RESULT " + json.dumps(result))
        sys.exit(1)

    venv_python = _get_venv_python(install_path)
    capabilities["python"] = str(venv_python)

    # Redirect pip's TEMP/TMP and cache dir onto the same drive as the install path.
    # Package/model downloads here run into the GB range; pip's default temp/cache
    # dirs live under the system TEMP (often a small C: drive) regardless of where
    # the venv lives, so a small system drive fills up and aborts the install even
    # when the target drive has plenty of free space.
    pip_env = _pip_env(install_path)

    # Upgrade pip/setuptools/wheel first — HY-World-2.0's dependency tree includes
    # packages using the newer PEP 639 `project.license = {text: ...}` pyproject.toml
    # format, which older vendored setuptools (bundled with fresh venvs) fails to parse
    # ("invalid pyproject.toml config: `project.license`"). Upgrading fixes this.
    log.info("─── UPGRADE PIP/SETUPTOOLS ───")
    try:
        subprocess.run(
            [str(venv_python), "-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel"],
            check=True,
            capture_output=False,
            env=pip_env,
        )
    except subprocess.CalledProcessError as e:
        log.warning("Could not upgrade pip/setuptools: %s (continuing anyway)", e)

    # Install base dependencies (requirements.txt only, not the world-gen pipeline)
    log.info("─── INSTALL DEPENDENCIES ───")
    if not _pip_install(venv_python, install_path, "requirements.txt", pip_env):
        result["code"] = "pip_failed"
        print("HYWORLD_INSTALL_RESULT " + json.dumps(result))
        sys.exit(1)

    # Best-effort: gsplat_maskgaussian. worldrecon's WorldMirror model class
    # unconditionally imports gsplat even when the "gs" output head is disabled,
    # so a stub is needed regardless (not just for GS-output users) — otherwise
    # importing the model at all fails on machines without a CUDA toolchain.
    if not args.skip_gsplat:
        log.info("─── OPTIONAL: gsplat_maskgaussian ───")
        gsplat_path = (
            install_path / "hyworld2" / "worldgen" / "third_party" / "gsplat_maskgaussian"
        )
        if _pip_install_editable(venv_python, gsplat_path, pip_env):
            capabilities["gsplat"] = True
            log.info("✓ gsplat_maskgaussian installed")
        else:
            log.warning("✗ gsplat_maskgaussian build failed (expected on Windows without build tools)")
            if _install_gsplat_stub(venv_python, pip_env):
                capabilities["gsplat"] = "stub"

    # Best-effort: flash-attn. Some of HY-World-2.0's modules hard-import flash_attn
    # with no further fallback (even though the accelerated kernels are only called
    # for fp16/bf16 tensors) — if the real package can't be built, install an
    # import-compatible stub instead so those modules still load on CPU/fp32.
    if not args.skip_flash_attn:
        log.info("─── OPTIONAL: flash-attn ───")
        if _pip_install_package(venv_python, "flash-attn", pip_env):
            capabilities["flash_attn"] = True
            log.info("✓ flash-attn installed")
        else:
            log.warning("✗ flash-attn build failed (expected on machines without CUDA toolchain)")
            if _install_flash_attn_stub(venv_python, pip_env):
                capabilities["flash_attn"] = "stub"

    # Best-effort: pybullet. Used by `zelpi pybullet load` to turn a HY-World
    # point cloud into a collidable physics environment. Ships prebuilt wheels
    # for most platforms; on machines without a matching wheel it builds from
    # source (slower, but no CUDA toolchain needed — pure CPU/C++ build).
    if not args.skip_pybullet:
        log.info("─── OPTIONAL: pybullet ───")
        if _pip_install_package(venv_python, "pybullet", pip_env):
            capabilities["pybullet"] = True
            log.info("✓ pybullet installed")
        else:
            capabilities["pybullet"] = False
            log.warning("✗ pybullet install failed")

    # Write capabilities manifest
    log.info("─── WRITE MANIFEST ───")
    manifest_path.write_text(json.dumps(capabilities, indent=2), encoding="utf-8")
    log.info("Manifest: %s", manifest_path)

    # Success
    log.info("─── SUCCESS ───")
    log.info("HY-World-2.0 installed at: %s", install_path)
    log.info("Python: %s", venv_python)
    log.info(
        "Capabilities: gsplat=%s, flash_attn=%s, pybullet=%s",
        capabilities["gsplat"], capabilities["flash_attn"], capabilities["pybullet"],
    )

    result["capabilities"] = capabilities
    print("HYWORLD_INSTALL_RESULT " + json.dumps(result))


if __name__ == "__main__":
    main()