#!/usr/bin/env python3
"""
Setup script for the Image Tagger application.
This script checks and installs all required dependencies.
"""

# Python 3.12+ compatibility patch for pkgutil.ImpImporter
import sys
if sys.version_info >= (3, 12):
    import pkgutil
    import importlib.machinery
    
    # Add ImpImporter as a compatibility shim for older packages
    if not hasattr(pkgutil, 'ImpImporter'):
        class ImpImporter:
            def __init__(self, path=None):
                self.path = path
            
            def find_module(self, fullname, path=None):
                return None
        
        pkgutil.ImpImporter = ImpImporter

import os
import sys
import subprocess
import platform
from pathlib import Path
import re
import urllib.request
import shutil
import tempfile
import time
import webbrowser

# Define the required packages
# Added setuptools and setuptools-distutils to fix distutils missing error
SETUPTOOLS_PACKAGES = [
    "setuptools>=58.0.0",
    "setuptools-distutils>=0.3.0",
    "wheel>=0.38.0",
]

REQUIRED_PACKAGES = [
    "streamlit>=1.21.0",
    "pillow>=9.0.0",
    # Important: Pin NumPy to 1.24.x for PyTorch compatibility
    # NumPy 2.x is not compatible with current PyTorch/torchvision builds
    "numpy==1.24.3",
    # Required for building Flash Attention
    "ninja>=1.10.0",
    "packaging>=20.0",
    # Required for the essence generator
    "matplotlib>=3.5.0",
    "tqdm>=4.62.0",
    "scipy>=1.7.0",
]

# Packages to install after PyTorch
POST_TORCH_PACKAGES = [
    "einops>=0.6.1",  # Required for Flash Attention
]

CUDA_PACKAGES = {
    # CUDA version: torch version
    "11.8": "torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118",
    "11.7": "torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117",
    "11.6": "torch==2.0.1+cu116 torchvision==0.15.2+cu116 --index-url https://download.pytorch.org/whl/cu116",
    "cpu": "torch==2.0.1+cpu torchvision==0.15.2+cpu --index-url https://download.pytorch.org/whl/cpu"
}

# ONNX and acceleration packages
ONNX_PACKAGES = [
    "onnx>=1.14.0",
    "onnxruntime>=1.15.0",
    "onnxruntime-gpu>=1.15.0;platform_system!='Darwin'",  # Skip GPU version on macOS
]

# TensorRT packages (for NVIDIA GPUs only)
TENSORRT_PACKAGES = [
    "nvidia-tensorrt>=8.6.1;platform_system=='Linux'",    # Linux only
    "tensorrt>=8.6.1;platform_system=='Windows'",         # Windows only
    "nvidia-cuda-nvrtc-cu11>=11.8.89;platform_system!='Darwin'",
    "nvidia-cuda-runtime-cu11>=11.8.89;platform_system!='Darwin'",
    "nvidia-cudnn-cu11>=8.7.0;platform_system!='Darwin'",
]

# Colors for terminal output
class Colors:
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'

def print_colored(text, color):
    """Print text in color"""
    if sys.platform == "win32":
        # Just print the text without color on Windows
        print(text)
    else:
        print(f"{color}{text}{Colors.ENDC}")

def check_python_version():
    """Check if Python version is 3.8 or higher, recommend 3.11.9 specifically"""
    print_colored("Checking Python version...", Colors.BLUE)
    
    version = sys.version_info
    if version.major < 3 or (version.major == 3 and version.minor < 8):
        print_colored("Error: Python 3.8 or higher is required. You have " + sys.version, Colors.FAIL)
        print_colored("Please install a newer Python version and try again.", Colors.FAIL)
        return False
    
    print_colored(f"[OK] Python {version.major}.{version.minor}.{version.micro} detected", Colors.GREEN)
    
    # Recommend Python 3.11.9 specifically
    if version.major == 3 and (version.minor != 11 or version.micro != 9):
        recommended = False
        warning_color = Colors.WARNING
        
        # Extra warning for Python 3.12+ due to known compatibility issues
        if version.major == 3 and version.minor >= 12:
            print_colored("WARNING: Python 3.12+ has known compatibility issues with this application.", Colors.FAIL)
            print_colored("The application has been tested and works reliably with Python 3.11.9.", Colors.FAIL)
            warning_color = Colors.FAIL
            recommended = True
        elif version.major == 3 and version.minor == 11 and version.micro != 9:
            print_colored("Note: This application has been tested with Python 3.11.9 specifically.", Colors.WARNING)
            recommended = True
        else:
            print_colored("Note: This application is recommended to use Python 3.11.9.", Colors.WARNING)
            recommended = True
            
        if recommended:
            print_colored("Download Python 3.11.9:", warning_color)
            if sys.platform == "win32":
                print_colored("  https://www.python.org/ftp/python/3.11.9/python-3.11.9-amd64.exe", warning_color)
            elif sys.platform == "darwin":  # macOS
                print_colored("  https://www.python.org/ftp/python/3.11.9/python-3.11.9-macos11.pkg", warning_color)
            else:  # Linux
                print_colored("  https://www.python.org/ftp/python/3.11.9/Python-3.11.9.tgz", warning_color)
                print_colored("  Or use your distribution's package manager", warning_color)
                
            if version.major == 3 and version.minor >= 12:
                # For Python 3.12+, we'll still continue but with a confirmation
                print_colored("\nDo you want to continue with the current Python version? (y/n)", Colors.BLUE)
                response = input().strip().lower()
                if response != 'y' and response != 'yes':
                    print_colored("Setup aborted. Please install Python 3.11.9 and try again.", Colors.FAIL)
                    return False
                print_colored("Continuing with current Python version. Some features may not work correctly.", Colors.WARNING)
    else:
        print_colored("[PERFECT] Python 3.11.9 detected - this is the recommended version!", Colors.GREEN)
    
    return True

def create_virtual_env():
    """Create a virtual environment if one doesn't exist"""
    print_colored("\nChecking for virtual environment...", Colors.BLUE)
    
    venv_path = Path("venv")
    if venv_path.exists():
        print_colored("[OK] Virtual environment already exists", Colors.GREEN)
        return True
    
    print_colored("Creating a new virtual environment...", Colors.BLUE)
    try:
        subprocess.run([sys.executable, "-m", "venv", "venv"], check=True)
        print_colored("[OK] Virtual environment created successfully", Colors.GREEN)
        return True
    except subprocess.CalledProcessError:
        print_colored("Error: Failed to create virtual environment", Colors.FAIL)
        return False

def get_venv_python():
    """Get path to Python in the virtual environment"""
    if sys.platform == "win32":
        return os.path.join("venv", "Scripts", "python.exe")
    else:
        return os.path.join("venv", "bin", "python")

def get_venv_pip():
    """Get path to pip in the virtual environment"""
    if sys.platform == "win32":
        return os.path.join("venv", "Scripts", "pip.exe")
    else:
        return os.path.join("venv", "bin", "pip")

def check_cuda():
    """Check CUDA availability and version"""
    print_colored("\nChecking for CUDA...", Colors.BLUE)
    
    # Check if nvidia-smi is available
    cuda_available = False
    cuda_version = None
    
    try:
        if sys.platform == "win32":
            # Windows: Check for nvidia-smi
            process = subprocess.run(["where", "nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if process.returncode == 0:
                # Run nvidia-smi to get version
                nvidia_smi = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                if nvidia_smi.returncode == 0:
                    cuda_available = True
                    # Extract CUDA Version from nvidia-smi output
                    match = re.search(r"CUDA Version: (\d+\.\d+)", nvidia_smi.stdout)
                    if match:
                        cuda_version = match.group(1)
        else:
            # Linux/Mac: Check for nvidia-smi
            process = subprocess.run(["which", "nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if process.returncode == 0:
                # Run nvidia-smi to get version
                nvidia_smi = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                if nvidia_smi.returncode == 0:
                    cuda_available = True
                    # Extract CUDA Version from nvidia-smi output
                    match = re.search(r"CUDA Version: (\d+\.\d+)", nvidia_smi.stdout)
                    if match:
                        cuda_version = match.group(1)
    except Exception as e:
        print_colored(f"Error checking CUDA: {str(e)}", Colors.WARNING)
    
    if cuda_available and cuda_version:
        print_colored(f"[OK] CUDA {cuda_version} detected", Colors.GREEN)
        # Return the closest supported CUDA version
        for supported_version in CUDA_PACKAGES.keys():
            if supported_version != "cpu" and float(supported_version) <= float(cuda_version):
                return supported_version
    
    print_colored("No CUDA detected, will use CPU-only version", Colors.WARNING)
    return "cpu"

def check_numpy_version():
    """Check if NumPy is installed and if it's a compatible version"""
    print_colored("\nChecking for existing NumPy installation...", Colors.BLUE)
    
    pip_path = get_venv_pip()
    
    # Check if NumPy is installed
    try:
        result = subprocess.run([pip_path, "show", "numpy"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode == 0:
            # Extract version from output
            version_match = re.search(r"Version: ([\d\.]+)", result.stdout)
            if version_match:
                current_version = version_match.group(1)
                print_colored(f"NumPy {current_version} is currently installed", Colors.BLUE)
                
                # Check if it's NumPy 2.x
                if current_version.startswith("2."):
                    print_colored(f"Warning: NumPy {current_version} is not compatible with PyTorch.", Colors.WARNING)
                    print_colored("NumPy will be downgraded to 1.24.3 for compatibility.", Colors.WARNING)
                    
                    # Uninstall NumPy
                    print_colored("Uninstalling incompatible NumPy version...", Colors.BLUE)
                    subprocess.run([pip_path, "uninstall", "-y", "numpy"], check=True)
                    print_colored("[OK] Successfully uninstalled NumPy", Colors.GREEN)
                    return False
                
                # If it's NumPy 1.24.x, we're good
                if current_version.startswith("1.24."):
                    print_colored(f"[OK] NumPy {current_version} is compatible with PyTorch", Colors.GREEN)
                    return True
                
                # For any other 1.x version, warn but continue
                print_colored(f"Note: NumPy will be updated to 1.24.3 for optimal compatibility", Colors.BLUE)
                return False
    except Exception as e:
        print_colored(f"Error checking NumPy: {str(e)}", Colors.WARNING)
    
    print_colored("NumPy is not installed", Colors.BLUE)
    return False
    
def install_packages(cuda_version):
    """Install required packages using pip"""
    print_colored("\nInstalling required packages...", Colors.BLUE)
    
    pip_path = get_venv_pip()
    
    # First, upgrade pip
    try:
        subprocess.run([pip_path, "install", "--upgrade", "pip"], check=True)
        print_colored("[OK] Pip upgraded successfully", Colors.GREEN)
    except subprocess.CalledProcessError:
        print_colored("Warning: Failed to upgrade pip", Colors.WARNING)
    
    # Install setuptools packages first to ensure distutils is available
    print_colored("\nInstalling setuptools and distutils...", Colors.BLUE)
    for package in SETUPTOOLS_PACKAGES:
        try:
            print_colored(f"Installing {package}...", Colors.BLUE)
            subprocess.run([pip_path, "install", package], check=True)
            print_colored(f"[OK] Installed {package}", Colors.GREEN)
        except subprocess.CalledProcessError as e:
            print_colored(f"Warning: Issue installing {package}: {e}", Colors.WARNING)
            print_colored("Continuing installation process...", Colors.BLUE)
    
    # Check NumPy version and install/upgrade if needed
    numpy_compatible = check_numpy_version()
    
    # Install base packages (except torch)
    for package in REQUIRED_PACKAGES:
        if numpy_compatible and package.startswith("numpy"):
            # Skip NumPy if already at compatible version
            print_colored(f"Skipping {package} (already installed at compatible version)", Colors.GREEN)
            continue
        
        try:
            print_colored(f"Installing {package}...", Colors.BLUE)
            subprocess.run([pip_path, "install", package], check=True)
            print_colored(f"[OK] Installed {package}", Colors.GREEN)
        except subprocess.CalledProcessError as e:
            print_colored(f"Error installing {package}: {e}", Colors.FAIL)
            return False
    
    # Install PyTorch with appropriate CUDA version
    print_colored(f"\nInstalling PyTorch {'with CUDA support' if cuda_version != 'cpu' else '(CPU version)'}...", Colors.BLUE)
    torch_command = CUDA_PACKAGES[cuda_version].split()
    try:
        subprocess.run([pip_path, "install"] + torch_command, check=True)
        print_colored("[OK] PyTorch installed successfully", Colors.GREEN)
    except subprocess.CalledProcessError as e:
        print_colored(f"Error installing PyTorch: {e}", Colors.FAIL)
        print_colored("You may need to manually install PyTorch from https://pytorch.org/", Colors.WARNING)
        return False
    
    # Install any packages that need to come after PyTorch
    for package in POST_TORCH_PACKAGES:
        try:
            subprocess.run([pip_path, "install", package], check=True)
            print_colored(f"[OK] Installed {package}", Colors.GREEN)
        except subprocess.CalledProcessError as e:
            print_colored(f"Error installing {package}: {e}", Colors.FAIL)
            return False
    
    return True

def install_acceleration_packages(cuda_version, install_tensorrt=False):
    """Install ONNX Runtime and TensorRT packages if CUDA is available"""
    print_colored("\nInstalling ONNX Runtime and acceleration packages...", Colors.BLUE)
    
    pip_path = get_venv_pip()
    
    # Choose either onnxruntime or onnxruntime-gpu, not both
    if cuda_version != "cpu":
        onnx_package = "onnxruntime-gpu>=1.15.0"
    else:
        onnx_package = "onnxruntime>=1.15.0"
    
    try:
        print_colored(f"Installing ONNX...", Colors.BLUE)
        subprocess.run([pip_path, "install", "onnx>=1.14.0"], check=True)
        print_colored(f"Installing {onnx_package}...", Colors.BLUE)
        subprocess.run([pip_path, "install", onnx_package], check=True)
        print_colored(f"[OK] ONNX packages installed", Colors.GREEN)
    except subprocess.CalledProcessError as e:
        print_colored(f"Warning: Issue installing ONNX packages: {e}", Colors.WARNING)
    
    # Install TensorRT packages only if explicitly requested
    if cuda_version != "cpu" and install_tensorrt:
        print_colored("\nAttempting to install TensorRT packages...", Colors.BLUE)
        print_colored("Note: TensorRT installation might require manual steps depending on your system", Colors.WARNING)
        
        for package in TENSORRT_PACKAGES:
            try:
                print_colored(f"Installing {package.split(';')[0]}...", Colors.BLUE)
                subprocess.run([pip_path, "install", package], check=True)
                print_colored(f"[OK] Installed {package.split(';')[0]}", Colors.GREEN)
            except subprocess.CalledProcessError as e:
                print_colored(f"Warning: Issue installing {package.split(';')[0]}: {e}", Colors.WARNING)
                print_colored("TensorRT may need to be installed manually.", Colors.WARNING)
    else:
        print_colored("Skipping TensorRT installation (no CUDA detected)", Colors.BLUE)
    
    return True

def run_application():
    """Run the application after setup"""
    print_colored("\nLaunching the application...", Colors.BLUE)
    
    # Check if app.py exists
    app_path = "app.py"
    if not os.path.exists(app_path):
        print_colored(f"Error: {app_path} not found. Can't launch application.", Colors.FAIL)
        return False
    
    # Get the path to streamlit in the virtual environment
    if sys.platform == "win32":
        streamlit_path = os.path.join("venv", "Scripts", "streamlit.exe")
    else:
        streamlit_path = os.path.join("venv", "bin", "streamlit")
    
    if not os.path.exists(streamlit_path):
        print_colored(f"Error: Streamlit not found at {streamlit_path}", Colors.FAIL)
        return False
    
    # Launch application
    try:
        print_colored("\nStarting Image Tagger Application...", Colors.GREEN)
        
        # Open browser first
        time.sleep(2)  # Give it a moment before starting the app
        
        # Run streamlit in a subprocess
        command = [streamlit_path, "run", app_path]
        subprocess.Popen(command)
        
        print_colored("[OK] Application launched successfully", Colors.GREEN)
        return True
    except Exception as e:
        print_colored(f"Error launching application: {e}", Colors.FAIL)
        return False

def main():
    """Main setup function"""
    print_colored("=" * 60, Colors.HEADER)
    print_colored("      Image Tagger - Setup Script", Colors.HEADER)
    print_colored("      (Recommended: Python 3.11.9)", Colors.HEADER)
    print_colored("=" * 60, Colors.HEADER)
    
    # Check Python version
    if not check_python_version():
        return False
    
    # Create virtual environment
    if not create_virtual_env():
        return False
    
    # Check CUDA version
    cuda_version = check_cuda()
    
    # Install packages
    if not install_packages(cuda_version):
        return False

    # Ask if the user wants TensorRT (only if CUDA is available)
    install_tensorrt = False
    if cuda_version != "cpu":
        print_colored("\nTensorRT packages add 2-3GB to the environment size but can improve performance.", Colors.BLUE)
        print_colored("Would you like to install TensorRT support? (y/n)", Colors.BLUE)
        response = input().strip().lower()
        install_tensorrt = (response == 'y' or response == 'yes')
    
    # Install acceleration packages with the user's preference
    if not install_acceleration_packages(cuda_version, install_tensorrt):
        print_colored("Warning: Some acceleration packages could not be installed", Colors.WARNING)
        print_colored("The application will still work, but performance may be reduced", Colors.WARNING)
    
    print_colored("\n" + "=" * 60, Colors.HEADER)
    print_colored("           Setup completed successfully!", Colors.GREEN)
    print_colored("=" * 60, Colors.HEADER)
    
    print_colored("\nTo run the application, use:", Colors.BLUE)
    if sys.platform == "win32":
        print_colored("  run_app.bat  or  python run_app.py", Colors.BOLD)
    else:
        print_colored("  ./run_app.py", Colors.BOLD)
    
    # Ask if the user wants to run the application now
    print_colored("\nWould you like to run the application now? (y/n)", Colors.BLUE)
    response = input().strip().lower()
    
    if response == 'y' or response == 'yes':
        run_application()
    else:
        print_colored("\nYou can run the application later using the commands above or running the batch scripts.", Colors.BLUE)
    
    return True

if __name__ == "__main__":
    success = main()
    if not success:
        sys.exit(1)