#!/bin/bash

# Define color codes
GREEN='\e[32m'
BLUE='\e[34m'
YELLOW='\e[33m'
NC='\e[0m' # No Color

# Function for consistent log messages
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_skip() { echo -e "${YELLOW}[SKIPPED]${NC} $1"; }

# Function to check if a command exists
command_exists() {
    command -v "$1" &> /dev/null
}

# Function to handle user input with Y/n prompt
get_user_confirmation() {
    local prompt="$1"
    local input
    log_info "$prompt (Y/n)"
    read -p "> " input
    [[ -z "$input" || "$input" =~ ^[Yy]$ ]]
}

# Function to setup SSH directory and keys
setup_ssh() {
    # Create .ssh directory if it doesn't exist
    if [ ! -d "$HOME/.ssh" ]; then
        mkdir -m 700 "$HOME/.ssh"
        log_success ".ssh-Verzeichnis wurde erstellt."
    else
        log_info ".ssh-Verzeichnis existiert bereits."
    fi

    # Ask for SSH public key
    log_info "SSH-Public-Key eingeben (Leerlassen zum Überspringen)"
    read -p "> " user_input

    if [ ! -z "$user_input" ]; then
        if grep -q "$user_input" "$HOME/.ssh/authorized_keys"; then
            log_info "Der SSH-Public-Key ist bereits vorhanden."
        else
            echo "$user_input" >> "$HOME/.ssh/authorized_keys"
            log_success "SSH-Public-Key wurde hinzugefügt."
        fi
    else
        log_skip "SSH-Key Installation übersprungen."
    fi
}

# Function to update system packages
update_system() {
    if get_user_confirmation "System aktualisieren?"; then
        log_info "System wird aktualisiert..."
        apt update && apt upgrade -y
        log_success "System wurde erfolgreich aktualisiert."
    else
        log_skip "System-Aktualisierung übersprungen."
    fi
}

# Function to install Docker and Dockge
install_docker() {
    if get_user_confirmation "Docker installieren?"; then
        if command_exists docker; then
            log_info "Docker ist bereits installiert."
        else
            log_info "Docker wird installiert..."
            # Install required packages
            apt install -y apt-transport-https ca-certificates curl software-properties-common
            
            # Add Docker's official GPG key
            curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
            
            # Add Docker repository
            echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
            
            # Update and install Docker
            apt update
            apt install -y docker-ce docker-ce-cli containerd.io
            
            # Start and enable Docker
            systemctl start docker
            systemctl enable docker
            log_success "Docker wurde erfolgreich installiert."
        fi

        # Install Dockge if Docker is available
        if command_exists docker; then
            if get_user_confirmation "Dockge (Docker Compose Manager) installieren?"; then
                log_info "Dockge wird installiert..."
                docker run -d --name=dockge -p 5001:5001 -v /var/run/docker.sock:/var/run/docker.sock louislam/dockge
                log_success "Dockge wurde erfolgreich installiert und läuft auf Port 5001."
            else
                log_skip "Dockge-Installation übersprungen."
            fi
        fi
    else
        log_skip "Docker-Installation übersprungen."
    fi
}

# Function to install Node.js
install_nodejs() {
    if get_user_confirmation "Node.js installieren?"; then
        if command_exists node; then
            log_info "Node.js ist bereits installiert."
        else
            log_info "Node.js wird installiert..."
            apt install -y curl
            curl -fsSL https://deb.nodesource.com/setup_lts.x | bash -
            apt install -y nodejs
            log_success "Node.js wurde erfolgreich installiert."
        fi
    else
        log_skip "Node.js-Installation übersprungen."
    fi
}

# Function to setup basic system tools and security
setup_basic_system() {
    if get_user_confirmation "Basic tools und Firewall installieren?"; then
        log_info "Installiere basic tools und UFW..."
        apt install -y htop zip unzip ufw
        
        # Configure UFW
        ufw allow ssh
        echo "y" | ufw enable
        
        log_success "Basic tools und UFW wurden erfolgreich installiert."
        log_success "SSH-Verbindungen sind erlaubt."
    else
        log_skip "Basic tools und UFW Installation übersprungen."
    fi
}

# Main execution
main() {
    setup_ssh
    update_system
    setup_basic_system    # Replace both previous functions
    install_docker
    install_nodejs
}

# Run main function
main