#!/bin/bash

# ComfyUI Launcher Installation Script
# This script automates the installation process for different Linux distributions

set -e  # Exit on any error

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored output
print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Detect Linux distribution
detect_distro() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        DISTRO=$ID
        VERSION=$VERSION_ID
    else
        print_error "Cannot detect Linux distribution"
        exit 1
    fi
}

# Install system dependencies based on distribution
install_system_deps() {
    print_status "Installing system dependencies for $DISTRO..."
    
    case $DISTRO in
        "fedora"|"rhel"|"centos"|"rocky")
            sudo dnf install -y cairo-devel gobject-introspection-devel gtk4-devel pkg-config python3-devel python3-pip
            ;;
        "ubuntu"|"debian")
            sudo apt update
            sudo apt install -y libcairo2-dev libgirepository1.0-dev libgtk-4-dev pkg-config python3-dev python3-pip python3-venv
            ;;
        "arch"|"manjaro")
            sudo pacman -S --needed cairo gobject-introspection gtk4 pkgconf python python-pip
            ;;
        *)
            print_warning "Unsupported distribution: $DISTRO"
            print_warning "Please install the following packages manually:"
            print_warning "- Cairo development libraries"
            print_warning "- GObject Introspection development libraries"
            print_warning "- GTK4 development libraries"
            print_warning "- pkg-config"
            print_warning "- Python 3.9+ development headers"
            ;;
    esac
}

# Check if conda is installed
check_conda() {
    if command -v conda &> /dev/null; then
        print_success "Conda found: $(conda --version)"
        return 0
    else
        print_warning "Conda not found"
        return 1
    fi
}

# Install conda if not present
install_conda() {
    print_status "Installing Miniconda..."
    
    # Detect architecture
    ARCH=$(uname -m)
    if [ "$ARCH" = "x86_64" ]; then
        CONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh"
    elif [ "$ARCH" = "aarch64" ]; then
        CONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh"
    else
        print_error "Unsupported architecture: $ARCH"
        exit 1
    fi
    
    # Download and install conda
    wget "$CONDA_URL" -O miniconda.sh
    bash miniconda.sh -b -p "$HOME/miniconda3"
    rm miniconda.sh
    
    # Add conda to PATH
    echo 'export PATH="$HOME/miniconda3/bin:$PATH"' >> ~/.bashrc
    export PATH="$HOME/miniconda3/bin:$PATH"
    
    # Initialize conda
    conda init bash
    
    print_success "Miniconda installed successfully"
    print_warning "Please restart your terminal or run: source ~/.bashrc"
}

# Setup Python virtual environment
setup_venv() {
    print_status "Setting up Python virtual environment..."
    
    if [ ! -d ".venv" ]; then
        python3 -m venv .venv
        print_success "Virtual environment created"
    else
        print_status "Virtual environment already exists"
    fi
    
    # Activate virtual environment
    source .venv/bin/activate
    
    # Upgrade pip
    pip install --upgrade pip
    
    # Install requirements
    if [ -f "requirements.txt" ]; then
        print_status "Installing Python dependencies..."
        pip install -r requirements.txt
        print_success "Python dependencies installed"
    else
        print_warning "requirements.txt not found, installing basic dependencies..."
        pip install PyGObject
    fi
}

# Make scripts executable
make_executable() {
    print_status "Making scripts executable..."
    chmod +x main.py 2>/dev/null || true
    chmod +x run.sh 2>/dev/null || true
    print_success "Scripts are now executable"
}

# Create desktop entry (optional)
create_desktop_entry() {
    if [ -f "comfyui-launcher.desktop" ]; then
        read -p "Install desktop entry? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            DESKTOP_FILE="$HOME/.local/share/applications/comfyui-launcher.desktop"
            mkdir -p "$(dirname "$DESKTOP_FILE")"
            
            # Update paths in desktop file
            CURRENT_DIR=$(pwd)
            sed "s|/home/enne2/Development/gtk-app|$CURRENT_DIR|g" comfyui-launcher.desktop > "$DESKTOP_FILE"
            
            print_success "Desktop entry installed to $DESKTOP_FILE"
        fi
    fi
}

# Main installation function
main() {
    print_status "ComfyUI Launcher Installation Script"
    print_status "======================================"
    
    # Detect distribution
    detect_distro
    print_status "Detected distribution: $DISTRO $VERSION"
    
    # Install system dependencies
    install_system_deps
    
    # Check for conda
    if ! check_conda; then
        read -p "Install Miniconda? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            install_conda
        else
            print_warning "Conda not installed. You'll need to install it manually for full functionality."
        fi
    fi
    
    # Setup Python environment
    setup_venv
    
    # Make scripts executable
    make_executable
    
    # Optional desktop entry
    create_desktop_entry
    
    print_success "Installation completed!"
    print_status "To run the launcher:"
    print_status "  ./run.sh"
    print_status "or:"
    print_status "  source .venv/bin/activate && python main.py"
}

# Run main function
main "$@"