#!/usr/bin/env python3
# fix_isotropic3d_deps.py
# Otomasi perbaikan libigl (igl) & (opsional) tinycudann untuk Isotropic3D/threestudio.

import os
import sys
import subprocess
import shlex
import json
import glob
from pathlib import Path

IS_ROOT = (os.geteuid() == 0) if hasattr(os, "geteuid") else False

def run(cmd, env=None, check=True):
    print(f"\n>>> {cmd}")
    result = subprocess.run(shlex.split(cmd), env=env, stdout=sys.stdout, stderr=sys.stderr)
    if check and result.returncode != 0:
        print(f"[ERROR] Command failed: {cmd}")
        sys.exit(result.returncode)
    return result.returncode

def py_run(code, check=True):
    print(f"\n>>> python - <<'PY' ...")
    p = subprocess.Popen([sys.executable, "-c", code])
    p.wait()
    if check and p.returncode != 0:
        print("[ERROR] Inline python failed.")
        sys.exit(p.returncode)
    return p.returncode

def detect_torch():
    code = r"""
import json, sys
try:
    import torch
    info = {
        "torch_version": torch.__version__,
        "cuda_runtime": getattr(torch.version, "cuda", None),
        "cuda_available": torch.cuda.is_available(),
        "device_capability": None,
        "device_name": None,
    }
    if info["cuda_available"]:
        info["device_name"] = torch.cuda.get_device_name(0)
        info["device_capability"] = torch.cuda.get_device_capability(0)
    print(json.dumps(info))
except Exception as e:
    print(json.dumps({"error": str(e)}))
"""
    out = subprocess.check_output([sys.executable, "-c", code], text=True)
    return json.loads(out.strip())

def set_cuda_env(cuda_home_guess=None):
    # best-effort set CUDA paths
    candidates = []
    if cuda_home_guess:
        candidates.append(cuda_home_guess)
    candidates += [
        "/usr/local/cuda",
        "/usr/local/cuda-12.5", "/usr/local/cuda-12.4", "/usr/local/cuda-12.3",
        "/usr/local/cuda-12.2", "/usr/local/cuda-12.1", "/usr/local/cuda-12.0",
        "/usr/local/cuda-11.8", "/usr/local/cuda-11.7",
    ]
    for c in candidates:
        if Path(c).exists():
            os.environ["CUDA_HOME"] = c
            os.environ["PATH"] = f"{c}/bin:" + os.environ.get("PATH", "")
            os.environ["LD_LIBRARY_PATH"] = f"{c}/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")
            print(f"[INFO] Using CUDA at {c}")
            return
    print("[WARN] CUDA not found on typical paths; proceeding anyway.")

def apt_install(packages):
    if not IS_ROOT:
        print("[WARN] Skipping apt-get because this process is not root.")
        return
    run("apt-get update")
    run(f"apt-get install -y {' '.join(packages)}")

def pip_install(pkgs, flags=""):
    run(f"{sys.executable} -m pip install -U {flags} {' '.join(pkgs)}")

def remove_old_igl():
    # pip uninstall
    run(f"{sys.executable} -m pip uninstall -y igl || true", check=False)
    # hard cleanup
    code = r"""
import sys, os, glob, shutil
removed = []
for sp in sys.path:
    for pat in ("igl", "igl-*"):
        for p in glob.glob(os.path.join(sp, pat)):
            try:
                if os.path.isdir(p):
                    shutil.rmtree(p, ignore_errors=True)
                else:
                    os.remove(p)
                removed.append(p)
            except Exception: pass
print("\n".join(removed))
"""
    print("[INFO] Removing leftover igl files (if any):")
    py_run(code, check=False)

def build_libigl(src_dir="/Isotropic3D/libigl-python-bindings"):
    # clone if needed
    if not Path(src_dir).exists():
        run(f"git clone --recursive https://github.com/libigl/libigl-python-bindings.git {src_dir}")
    else:
        run(f"git -C {src_dir} submodule update --init --recursive")

    # install with no build isolation so it sees our env
    run(f"{sys.executable} -m pip install -v --no-build-isolation .", env=os.environ | {}, check=True)

def verify_igl():
    code = r"""
import igl
print("igl module path:", igl.__file__)
print("has fast_winding_number_for_meshes:", hasattr(igl, "fast_winding_number_for_meshes"))
"""
    py_run(code)

def compute_arch_from_capability(cap=None):
    # map major.minor to CMake arch (rounded typical)
    # Lovelace(8.9)=89, Ampere(8.6)=86, A100(8.0)=80, Turing(7.5)=75, Volta(7.0)=70
    mapping = {(8,9): 89, (8,6): 86, (8,0): 80, (7,5): 75, (7,0): 70}
    if isinstance(cap, (list, tuple)) and len(cap) >= 2:
        return mapping.get((cap[0], cap[1]), cap[0]*10 + cap[1])
    return None

def ensure_tinycudann(src_dir="/Isotropic3D/tiny-cuda-nn", arch=None):
    # uninstall old
    run(f"{sys.executable} -m pip uninstall -y tinycudann tiny-cuda-nn || true", check=False)

    # clone if needed
    if not Path(src_dir).exists():
        run(f"git clone --recursive https://github.com/NVlabs/tiny-cuda-nn.git {src_dir}")
    else:
        run(f"git -C {src_dir} submodule update --init --recursive")

    bind_dir = Path(src_dir) / "bindings" / "torch"
    if not bind_dir.exists():
        print("[ERROR] tiny-cuda-nn bindings/torch not found.")
        sys.exit(2)

    env = os.environ.copy()
    if arch:
        env["TCNN_CUDA_ARCHITECTURES"] = str(arch)
        env["TORCH_CUDA_LIST"] = ""  # avoid confusion
    # Build with no build isolation so it can use our Torch/pybind11
    run(f"{sys.executable} -m pip install -v --no-build-isolation .", env=env | {}, check=True)

def verify_tcnn():
    code = r"""
import torch
import tinycudann as tcnn
print("tinycudann import OK; torch", torch.__version__)
"""
    py_run(code)

def main():
    import argparse
    ap = argparse.ArgumentParser(description="Fix igl and (optional) tinycudann for Isotropic3D.")
    ap.add_argument("--cuda-home", type=str, default=None, help="Custom CUDA_HOME path (optional).")
    ap.add_argument("--do-tcnn", action="store_true", help="Also build tiny-cuda-nn bindings.")
    ap.add_argument("--libigl-src", type=str, default="/Isotropic3D/libigl-python-bindings")
    ap.add_argument("--tcnn-src", type=str, default="/Isotropic3D/tiny-cuda-nn")
    args = ap.parse_args()

    # 0) Toolchain + backend
    apt_install(["build-essential", "cmake", "ninja-build", "git", "libeigen3-dev"])
    pip_install(["pip", "wheel", "setuptools", "scikit-build-core", "pybind11", "numpy"])

    # 1) CUDA env (best effort)
    set_cuda_env(args.cuda_home)

    # 2) Reinstall igl properly
    print("\n=== Reinstalling libigl (igl) ===")
    remove_old_igl()
    # Install from local clone or clone to default dir then install
    if Path(args.libigl_src).exists():
        run(f"git -C {args.libigl_src} submodule update --init --recursive")
        run(f"{sys.executable} -m pip install -v --no-build-isolation .", env=os.environ | {}, check=True)
    else:
        build_libigl(args.libigl_src)
    print("\n=== Verifying igl ===")
    verify_igl()

    # 3) (Optional) Build tinycudann
    if args.do_tcnn:
        print("\n=== Building tiny-cuda-nn / tinycudann ===")
        info = detect_torch()
        print(f"[INFO] Torch info: {info}")
        arch = None
        if info.get("cuda_available") and info.get("device_capability"):
            arch = compute_arch_from_capability(info["device_capability"])
            print(f"[INFO] Detected GPU capability {info['device_capability']} -> arch {arch}")
        else:
            print("[WARN] CUDA not available from torch; proceeding without arch hint.")
        ensure_tinycudann(args.tcnn_src, arch=arch)
        print("\n=== Verifying tinycudann ===")
        verify_tcnn()

    print("\n[OK] All done. You can now run Isotropic3D:")
    print("    cd /Isotropic3D")
    print("    python launch.py --config configs/isotropic3d-shading.yaml --train --gpu 0")

if __name__ == "__main__":
    main()