#!/usr/bin/env python3
"""
Auto-install script for ComfyUI RIFE TensorRT Auto.
Detects the local CUDA version and installs only the matching TensorRT wheels.
The NVIDIA CUDA Toolkit must already be installed on the system.
"""

import subprocess
import sys
import os
import re
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent


def run_command(cmd):
    """Run a command (argv list) and return stdout + return code."""
    try:
        return subprocess.check_output(cmd, stderr=subprocess.DEVNULL, text=True).strip(), 0
    except subprocess.CalledProcessError as e:
        return "", e.returncode
    except Exception:
        return "", 1


def detect_cuda_version():
    """Detect CUDA version from nvcc, CUDA_PATH or CUDA_HOME."""
    print("Detecting CUDA version...")

    # Try nvcc command
    stdout, returncode = run_command(["nvcc", "--version"])
    if returncode == 0:
        match = re.search(r"release (\d+\.\d+)", stdout)
        if match:
            version = match.group(1)
            print(f"Found CUDA version: {version}")
            return version

    # Try CUDA_PATH
    cuda_path = os.getenv("CUDA_PATH")
    if cuda_path:
        nvcc_path = os.path.join(cuda_path, "bin", "nvcc")
        if os.path.exists(nvcc_path):
            stdout, returncode = run_command([nvcc_path, "--version"])
            if returncode == 0:
                match = re.search(r"release (\d+\.\d+)", stdout)
                if match:
                    version = match.group(1)
                    print(f"Found CUDA version via CUDA_PATH: {version}")
                    return version

    # Try CUDA_HOME
    cuda_home = os.getenv("CUDA_HOME")
    if cuda_home:
        nvcc_path = os.path.join(cuda_home, "bin", "nvcc")
        if os.path.exists(nvcc_path):
            stdout, returncode = run_command([nvcc_path, "--version"])
            if returncode == 0:
                match = re.search(r"release (\d+\.\d+)", stdout)
                if match:
                    version = match.group(1)
                    print(f"Found CUDA version via CUDA_HOME: {version}")
                    return version

    print("Could not detect CUDA version automatically.")
    print("Please ensure the NVIDIA CUDA Toolkit is installed and nvcc is in your PATH,")
    print("or set CUDA_PATH / CUDA_HOME.")
    return None


def install_requirements(cuda_version):
    """Install the correct TensorRT packages for the detected CUDA version."""
    major_version = int(cuda_version.split('.')[0])
    print(f"Installing requirements for CUDA {major_version}...")

    if major_version == 13:
        print("Using CUDA 13 TensorRT packages (RTX 50 series)")
        req_file = "requirements_cu13.txt"
    elif major_version == 12:
        print("Using CUDA 12 TensorRT packages (RTX 30/40 series)")
        req_file = "requirements_cu12.txt"
    else:
        print(f"Unsupported CUDA version: {cuda_version}")
        print("Supported versions: CUDA 12.x, CUDA 13.x")
        return False

    req_path = SCRIPT_DIR / req_file
    if not req_path.exists():
        print(f"Missing requirements file: {req_path}")
        return False

    # Install base dependencies first, then the CUDA-specific TensorRT wheels.
    # Using a single pip call per file lets pip resolve and cache dependencies efficiently.
    for req_name in ["requirements.txt", req_file]:
        req_file_path = SCRIPT_DIR / req_name
        if not req_file_path.exists():
            continue
        print(f"Installing from {req_name}...")
        try:
            subprocess.check_call(
                [sys.executable, "-m", "pip", "install", "--no-input", "--prefer-binary", "-r", str(req_file_path)]
            )
        except subprocess.CalledProcessError:
            print(f"Failed to install {req_name}")
            return False

    print("Installation completed successfully!")
    print("You can now use the ComfyUI RIFE TensorRT Auto node.")
    return True


def main():
    cuda_version = detect_cuda_version()
    if not cuda_version:
        sys.exit(1)

    success = install_requirements(cuda_version)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()