"""
Installation script for ComfyUI-TRELLIS2_Motion dependencies.

Run with: python install.py
For full installation including TRELLIS: python install.py --full
For Colab: python install.py --colab
"""

import subprocess
import sys
import os


def run_pip(*args, ignore_errors=False):
    """Run pip with the given arguments."""
    try:
        subprocess.check_call([sys.executable, "-m", "pip", *args])
        return True
    except subprocess.CalledProcessError as e:
        if ignore_errors:
            print(f"  Warning: pip command failed: {e}")
            return False
        raise


def run_command(cmd, ignore_errors=False):
    """Run a shell command."""
    try:
        subprocess.check_call(cmd, shell=True)
        return True
    except subprocess.CalledProcessError as e:
        if ignore_errors:
            print(f"  Warning: command failed: {e}")
            return False
        raise


def install_dependencies(full=False, colab=False):
    """Install all required dependencies."""
    print("="*60)
    print("Installing ComfyUI-TRELLIS2_Motion dependencies...")
    print("="*60)

    # Core dependencies (usually pre-installed)
    print("\n[1/7] Checking core dependencies...")
    run_pip("install", "--upgrade", "numpy>=1.21.0", "pillow>=9.0.0", ignore_errors=True)

    # HuggingFace and transformers
    print("\n[2/7] Installing HuggingFace dependencies...")
    run_pip("install", "huggingface_hub>=0.20.0", "transformers>=4.30.0", "safetensors")

    # Mesh processing
    print("\n[3/7] Installing mesh processing libraries...")
    run_pip("install", "trimesh>=3.0.0")
    run_pip("install", "xatlas", ignore_errors=True)

    # Background removal
    print("\n[4/7] Installing background removal (rembg)...")
    run_pip("install", "rembg[gpu]" if colab else "rembg")

    # Image processing
    print("\n[5/7] Installing image processing libraries...")
    run_pip("install", "kornia>=0.7.0", "scipy>=1.10.0", "einops>=0.6.0")

    # nvdiffrast (CUDA mesh rendering)
    print("\n[6/7] Installing nvdiffrast (CUDA mesh rendering)...")
    if not run_pip("install", "nvdiffrast", ignore_errors=True):
        print("  Trying from git...")
        run_pip("install", "git+https://github.com/NVlabs/nvdiffrast.git", ignore_errors=True)

    # TRELLIS (only in full/colab mode)
    if full or colab:
        print("\n[7/7] Installing TRELLIS2...")

        # Install TRELLIS dependencies first
        run_pip("install", "omegaconf", "imageio", "imageio-ffmpeg", ignore_errors=True)

        # Try to install spconv for CUDA
        if colab:
            # Colab typically uses CUDA 12.x
            if not run_pip("install", "spconv-cu120", ignore_errors=True):
                run_pip("install", "spconv-cu118", ignore_errors=True)

        # Install TRELLIS from GitHub
        print("  Installing TRELLIS from GitHub (this may take a while)...")
        run_pip("install", "git+https://github.com/microsoft/TRELLIS.git", ignore_errors=True)
    else:
        print("\n[7/7] Skipping TRELLIS installation (use --full or --colab to install)")

    print("\n" + "="*60)
    print("Installation complete!")
    print("="*60)

    if not (full or colab):
        print("\nNOTE: TRELLIS2 was NOT installed.")
        print("To install TRELLIS2, run one of:")
        print("  python install.py --full    # Full installation")
        print("  python install.py --colab   # Optimized for Colab")
        print("\nOr install manually:")
        print("  pip install git+https://github.com/microsoft/TRELLIS.git")


def check_dependencies():
    """Check which dependencies are installed."""
    deps = {
        "torch": "torch",
        "numpy": "numpy",
        "PIL": "pillow",
        "huggingface_hub": "huggingface_hub",
        "transformers": "transformers",
        "trimesh": "trimesh",
        "kornia": "kornia",
        "cv2": "opencv-python",
        "nvdiffrast": "nvdiffrast",
        "rembg": "rembg",
    }

    print("Checking dependencies...")
    print("-" * 40)

    for import_name, package_name in deps.items():
        try:
            __import__(import_name)
            print(f"  [OK] {package_name}")
        except ImportError:
            print(f"  [MISSING] {package_name}")

    # Check CUDA
    try:
        import torch
        if torch.cuda.is_available():
            print(f"\n  CUDA: Available ({torch.cuda.get_device_name(0)})")
            print(f"  CUDA Version: {torch.version.cuda}")
        else:
            print("\n  CUDA: Not available (CPU only)")
    except Exception as e:
        print(f"\n  CUDA check failed: {e}")


if __name__ == "__main__":
    args = sys.argv[1:]

    if "check" in args:
        check_dependencies()
    elif "--help" in args or "-h" in args:
        print(__doc__)
        print("\nOptions:")
        print("  --full    Install all dependencies including TRELLIS2")
        print("  --colab   Install for Google Colab (includes TRELLIS2)")
        print("  check     Check which dependencies are installed")
    else:
        full = "--full" in args
        colab = "--colab" in args
        install_dependencies(full=full, colab=colab)
        print("\nRun 'python install.py check' to verify installation.")