# ==========================================================
# FILE: install.py
# ==========================================================
#!/usr/bin/env python3
import os, sys, subprocess, shutil, platform, venv

ENV_DIR   = "venv"
CACHE_DIR = "hf_download"
MODEL_ID  = "hunyuanvideo-community/HunyuanVideo"

def run(cmd: list[str]) -> None:
    print("❯", *cmd, flush=True)
    subprocess.check_call(cmd)

def ensure_python() -> None:
    if sys.version_info < (3, 8):
        sys.exit("Python 3.8 or newer required.")

def cuda_info() -> None:
    if shutil.which("nvidia-smi"):
        run(["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv"])
    else:
        print("⚠️  CUDA GPU not detected – falling back to CPU (slow).")

def create_venv() -> tuple[str, str]:
    if not os.path.isdir(ENV_DIR):
        print("📦 Creating virtual environment…")
        venv.create(ENV_DIR, with_pip=True)
    if platform.system() == "Windows":
        pip = os.path.join(ENV_DIR, "Scripts", "pip.exe")
        py  = os.path.join(ENV_DIR, "Scripts", "python.exe")
    else:
        pip = os.path.join(ENV_DIR, "bin", "pip")
        py  = os.path.join(ENV_DIR, "bin", "python")
    return pip, py

def install_reqs(pip: str) -> None:
    run([pip, "install", "--upgrade", "pip"])
    run([pip, "install", "-r", "requirements.txt"])

def hf_login(py: str) -> None:
    run([py, "-c", "from diffusers_helper.hf_login import login; login()"])

def cache_model(py: str) -> None:
    os.makedirs(CACHE_DIR, exist_ok=True)
    code = (
        f"from huggingface_hub import snapshot_download;"
        f"snapshot_download('{MODEL_ID}', cache_dir='{CACHE_DIR}', resume_download=True)"
    )
    run([py, "-c", code])

def main() -> None:
    ensure_python()
    cuda_info()
    pip, py = create_venv()
    install_reqs(pip)
    hf_login(py)
    cache_model(py)
    if platform.system() == "Windows":
        print(f"\nActivate: {ENV_DIR}\\Scripts\\Activate.ps1")
    else:
        print(f"\nActivate: source {ENV_DIR}/bin/activate")
    print("Run:      python ghostpack.py --inbrowser")

if __name__ == "__main__":
    main()