#!/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",
    "ninja>=1.10.0",
    "packaging>=20.0",
    "matplotlib>=3.5.0",
    "tqdm>=4.62.0",
    "scipy>=1.7.0",
    "safetensors>=0.3.0",
]


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


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 ensure_compatible_numpy() -> bool:
    """
    Install a NumPy version compatible with the current Python version.

    - Python < 3.13: NumPy 1.24.x–1.x (we keep <2 to avoid ecosystem surprises)
    - Python >= 3.13: NumPy 2.1+ (since 1.26 wheels don't support 3.13)
    """
    print_colored("\nEnsuring NumPy compatibility...", Colors.BLUE)

    pip_path = get_venv_pip()

    if sys.version_info >= (3, 13):
        desired_spec = "numpy>=2.1.0,<3.0.0"
        desired_major = 2
    else:
        desired_spec = "numpy>=1.24.0,<2.0.0"
        desired_major = 1

    def parse_ver(text: str):
        m = re.search(r"Version:\s*([0-9]+)\.([0-9]+)\.([0-9]+)", text)
        if not m:
            return None
        return (int(m.group(1)), int(m.group(2)), int(m.group(3)))

    try:
        show = subprocess.run(
            [pip_path, "show", "numpy"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )

        if show.returncode == 0:
            ver = parse_ver(show.stdout)
            if ver:
                major, minor, patch = ver
                print_colored(f"Found NumPy {major}.{minor}.{patch}", Colors.BLUE)

                ok = (major == desired_major) and (major != 1 or minor >= 24)
                if ok:
                    print_colored("[OK] NumPy version is compatible", Colors.GREEN)
                    return False

                print_colored("NumPy version not compatible for this Python; reinstalling...", Colors.BLUE)
                subprocess.run([pip_path, "uninstall", "-y", "numpy"], check=False)

        print_colored(f"Installing {desired_spec}...", Colors.BLUE)
        subprocess.run([pip_path, "install", desired_spec], check=True)
        print_colored(f"[OK] Installed {desired_spec}", Colors.GREEN)
        return True

    except subprocess.CalledProcessError as e:
        print_colored(f"Error installing NumPy: {e}", Colors.FAIL)
        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  = ensure_compatible_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
    
    # Install PyTorch with appropriate CUDA version
    print_colored(
        f"\nInstalling PyTorch ({'GPU '+cuda_version if cuda_version != 'cpu' else 'CPU-only'})...",
        Colors.BLUE
    )
    if not install_pytorch(cuda_version):
        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 import successfully", Colors.GREEN)
        else:
            print_colored(f"Warning: Import test failed:\n{result.stderr}", Colors.WARNING)
            
    except Exception as e:
        print_colored(f"Warning: Could not perform compatibility check: {e}", Colors.WARNING)
    
    return True

def check_python_version():
    v = sys.version_info
    if (v.major, v.minor) < (3, 10) or (v.major, v.minor) >= (3, 14):
        print_colored(
            f"Error: Python 3.10–3.13 required. You have {v.major}.{v.minor}.",
            Colors.FAIL
        )
        return False
    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)
        return _pick_pytorch_tag_from_driver(cuda_version)

    print_colored("No CUDA detected, using CPU-only version", Colors.WARNING)
    return "cpu"


# Order matters: we try the newest runtime your driver supports, then fall back.
PYTORCH_CUDA_TAGS = [
    (13.0, "cu130"),
    (12.9, "cu129"),
    (12.8, "cu128"),
    (12.6, "cu126"),
    (12.4, "cu124"),
    (12.1, "cu121"),
    (11.8, "cu118"),
]

def _pick_pytorch_tag_from_driver(cuda_version_str: str) -> str:
    try:
        v = float(cuda_version_str)
    except Exception:
        return "cpu"

    for min_v, tag in PYTORCH_CUDA_TAGS:
        if v >= min_v:
            return tag
    return "cpu"

def _torch_install_attempt(pip_path: str, index_url: str, pre: bool = False) -> bool:
    # Always add PyPI as extra index so dependencies are resolvable.
    cmd = [pip_path, "install", "--upgrade"]
    if pre:
        cmd.append("--pre")
    cmd += [
        "torch", "torchvision",
        "--index-url", index_url,
        "--extra-index-url", "https://pypi.org/simple",
    ]
    try:
        subprocess.run(cmd, check=True)
        return True
    except subprocess.CalledProcessError:
        return False

def install_pytorch(cuda_tag: str) -> bool:
    """
    Install torch/torchvision.
    - On macOS: install from PyPI (MPS/CPU handled by upstream)
    - On Windows/Linux: install from the PyTorch CUDA/CPU index matching the driver.
    - If Python 3.13+ and stable fails, try nightly as a fallback.
    """
    pip_path = get_venv_pip()

    print_colored("\nInstalling PyTorch...", Colors.BLUE)

    # Clean out any prior torch installs (helps reruns / partial installs)
    subprocess.run([pip_path, "uninstall", "-y", "torch", "torchvision", "torchaudio"], check=False)

    if sys.platform == "darwin":
        try:
            subprocess.run([pip_path, "install", "--upgrade", "torch", "torchvision"], check=True)
            print_colored("[OK] PyTorch installed (macOS)", Colors.GREEN)
            return True
        except subprocess.CalledProcessError as e:
            print_colored(f"Error installing PyTorch on macOS: {e}", Colors.FAIL)
            return False

    # Build a fallback list from the detected tag downwards, then CPU last.
    all_tags = [tag for _, tag in PYTORCH_CUDA_TAGS]
    if cuda_tag == "cpu":
        candidates = ["cpu"]
    else:
        if cuda_tag in all_tags:
            start = all_tags.index(cuda_tag)
            candidates = all_tags[start:] + ["cpu"]
        else:
            candidates = ["cu118", "cpu"]

    # 1) Try stable
    for tag in candidates:
        index_url = f"https://download.pytorch.org/whl/{tag}"
        print_colored(f"Trying stable PyTorch from {index_url} ...", Colors.BLUE)
        if _torch_install_attempt(pip_path, index_url=index_url, pre=False):
            print_colored(f"[OK] PyTorch installed successfully ({tag})", Colors.GREEN)
            return True

    # 2) If on 3.13+, try nightly (optional but helps early adopters)
    if sys.version_info >= (3, 13):
        for tag in candidates:
            index_url = f"https://download.pytorch.org/whl/nightly/{tag}"
            print_colored(f"Trying nightly PyTorch from {index_url} ...", Colors.BLUE)
            if _torch_install_attempt(pip_path, index_url=index_url, pre=True):
                print_colored(f"[OK] PyTorch nightly installed successfully ({tag})", Colors.GREEN)
                return True

    print_colored(
        "Error: Could not install a compatible PyTorch build for this Python/CUDA combo.\n"
        "Tip: Python 3.12 is the most broadly supported choice if you hit this.",
        Colors.FAIL
    )
    return False

def _try_pip_install(pip_path: str, spec: str) -> bool:
    try:
        subprocess.run([pip_path, "install", spec], check=True)
        return True
    except subprocess.CalledProcessError:
        return False

def install_onnx_packages(cuda_version):
    """Install ONNX + an available ONNX Runtime for this OS/Python."""
    print_colored("\nInstalling ONNX packages...", Colors.BLUE)
    pip_path = get_venv_pip()

    # ONNX itself is fine
    if not _try_pip_install(pip_path, "onnx>=1.14.0"):
        print_colored("Error: failed to install onnx", Colors.FAIL)
        return False

    # Python 3.14+ currently often has no official onnxruntime* wheels on PyPI
    # (PyPI classifiers are up to 3.13 for onnxruntime/onnxruntime-gpu).
    if sys.version_info >= (3, 14):
        print_colored(
            "Warning: ONNX Runtime wheels are typically not available on PyPI for Python 3.14+ yet.\n"
            "ONNX installed, but runtime install is skipped. Use Python 3.12/3.13 for ONNX Runtime.",
            Colors.WARNING
        )
        
    # Decide what to try for runtime
    is_windows = (sys.platform == "win32")
    has_cuda = (cuda_version != "cpu")  # in your script this means NVIDIA/CUDA detected

    # Preferred order:
    # - If CUDA detected: onnxruntime-gpu
    # - On Windows without CUDA: onnxruntime-directml (GPU via DirectML)
    # - Fallback: onnxruntime (CPU)
    tried = []

    if has_cuda:
        tried.append("onnxruntime-gpu>=1.15.0")
        if _try_pip_install(pip_path, tried[-1]):
            print_colored("[OK] Installed onnxruntime-gpu", Colors.GREEN)
            return True

    if is_windows:
        tried.append("onnxruntime-directml>=1.15.0")
        if _try_pip_install(pip_path, tried[-1]):
            print_colored("[OK] Installed onnxruntime-directml", Colors.GREEN)
            return True

    tried.append("onnxruntime>=1.15.0")
    if _try_pip_install(pip_path, tried[-1]):
        print_colored("[OK] Installed onnxruntime (CPU)", Colors.GREEN)
        return True

    print_colored(
        "Warning: Could not install any ONNX Runtime variant.\n"
        f"Tried: {', '.join(tried)}",
        Colors.WARNING
    )
    return False

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)