#!/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
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",
    # CRITICAL: Pin NumPy to 1.24.x and prevent 2.x installation
    "numpy>=1.24.0,<2.0.0",
    "ninja>=1.10.0",
    "packaging>=20.0",
    "matplotlib>=3.5.0",
    "tqdm>=4.62.0",
    "scipy>=1.7.0",
    "safetensors>=0.3.0",
    "timm>=0.9.0",  # Add this line for PyTorch Image Models
]

# Packages to install after PyTorch
POST_TORCH_PACKAGES = [
    "einops>=0.6.1",
]

CUDA_PACKAGES = {
    "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'",
]

# 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":
        print(text)
    else:
        print(f"{color}{text}{Colors.ENDC}")

def check_and_fix_numpy():
    """Check for NumPy 2.x and fix compatibility issues"""
    print_colored("\nChecking NumPy compatibility...", Colors.BLUE)
    
    pip_path = get_venv_pip()
    
    try:
        result = subprocess.run([pip_path, "show", "numpy"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode == 0:
            version_match = re.search(r"Version: ([\d\.]+)", result.stdout)
            if version_match:
                current_version = version_match.group(1)
                print_colored(f"Found NumPy {current_version}", Colors.BLUE)
                
                if current_version.startswith("2."):
                    print_colored(f"ERROR: NumPy {current_version} is incompatible with PyTorch!", Colors.FAIL)
                    print_colored("Uninstalling NumPy 2.x and installing compatible version...", Colors.BLUE)
                    
                    # Force uninstall NumPy 2.x
                    subprocess.run([pip_path, "uninstall", "-y", "numpy"], check=True)
                    
                    # Install compatible NumPy version with constraints
                    subprocess.run([pip_path, "install", "numpy>=1.24.0,<2.0.0"], check=True)
                    
                    print_colored("[OK] NumPy downgraded to compatible version", Colors.GREEN)
                    return True
                    
                elif current_version.startswith("1.24."):
                    print_colored(f"[OK] NumPy {current_version} is compatible", Colors.GREEN)
                    return False
                else:
                    print_colored(f"Updating NumPy to recommended version...", Colors.BLUE)
                    subprocess.run([pip_path, "install", "numpy>=1.24.0,<2.0.0"], check=True)
                    return True
    except Exception as e:
        print_colored(f"Error checking NumPy: {e}", Colors.WARNING)
    
    return False

def install_packages(cuda_version):
    """Install required packages using pip"""
    print_colored("\nInstalling required packages...", Colors.BLUE)
    
    pip_path = get_venv_pip()
    
    # Upgrade pip first
    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
    print_colored("\nInstalling setuptools...", Colors.BLUE)
    for package in SETUPTOOLS_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"Warning: Issue installing {package}: {e}", Colors.WARNING)
    
    # Check and fix NumPy compatibility before installing other packages
    numpy_was_updated = check_and_fix_numpy()
    
    # Install base packages
    for package in REQUIRED_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"Error installing {package}: {e}", Colors.FAIL)
            return False
    
    # If NumPy was updated, we need to reinstall PyTorch to ensure compatibility
    if numpy_was_updated:
        print_colored("\nNumPy was updated, ensuring PyTorch compatibility...", Colors.BLUE)
        try:
            # Uninstall existing PyTorch
            subprocess.run([pip_path, "uninstall", "-y", "torch", "torchvision"], check=False)
        except:
            pass  # Ignore errors if not installed
    
    # 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)
        return False
    
    # Install post-PyTorch packages
    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
    
    # Final NumPy compatibility check
    print_colored("\nPerforming final compatibility check...", Colors.BLUE)
    try:
        # Test import in the virtual environment
        python_path = get_venv_python()
        test_cmd = [python_path, "-c", "import torch; import torchvision; import numpy; print('All imports successful')"]
        result = subprocess.run(test_cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print_colored("[OK] All packages are compatible", Colors.GREEN)
        else:
            print_colored(f"Warning: Compatibility test failed: {result.stderr}", Colors.WARNING)
            # Try to fix by reinstalling with --force-reinstall
            print_colored("Attempting to fix with force reinstall...", Colors.BLUE)
            subprocess.run([pip_path, "install", "--force-reinstall", "numpy>=1.24.0,<2.0.0"], check=True)
            
    except Exception as e:
        print_colored(f"Warning: Could not perform compatibility check: {e}", Colors.WARNING)
    
    return True

def check_python_version():
    """Check if Python version is 3.8 or higher"""
    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)
        return False
    
    print_colored(f"[OK] Python {version.major}.{version.minor}.{version.micro} detected", 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)
    
    cuda_available = False
    cuda_version = None
    
    try:
        if sys.platform == "win32":
            process = subprocess.run(["where", "nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        else:
            process = subprocess.run(["which", "nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        if process.returncode == 0:
            nvidia_smi = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if nvidia_smi.returncode == 0:
                cuda_available = True
                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)
        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, using CPU-only version", Colors.WARNING)
    return "cpu"

def install_onnx_packages(cuda_version):
    """Install ONNX packages"""
    print_colored("\nInstalling ONNX packages...", Colors.BLUE)
    
    pip_path = get_venv_pip()
    
    try:
        subprocess.run([pip_path, "install", "onnx>=1.14.0"], check=True)
        
        if cuda_version != "cpu":
            subprocess.run([pip_path, "install", "onnxruntime-gpu>=1.15.0"], check=True)
        else:
            subprocess.run([pip_path, "install", "onnxruntime>=1.15.0"], check=True)
            
        print_colored("[OK] ONNX packages installed", Colors.GREEN)
    except subprocess.CalledProcessError as e:
        print_colored(f"Warning: ONNX installation issues: {e}", Colors.WARNING)
    
    return True

def main():
    """Main setup function"""
    print_colored("=" * 60, Colors.HEADER)
    print_colored("      Image Tagger - Setup Script", Colors.HEADER)
    print_colored("=" * 60, Colors.HEADER)
    
    if not check_python_version():
        return False
    
    if not create_virtual_env():
        return False
    
    cuda_version = check_cuda()
    
    if not install_packages(cuda_version):
        return False

    if not install_onnx_packages(cuda_version):
        print_colored("Warning: ONNX packages had issues", Colors.WARNING)
    
    print_colored("\n" + "=" * 60, Colors.HEADER)
    print_colored("           Setup completed successfully!", Colors.GREEN)
    print_colored("=" * 60, Colors.HEADER)
    
    return True

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