#!/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
        log_info "Node.js wird installiert/aktualisiert..."
        
        # Install npm first
        apt install -y npm
        
        # Update npm to latest version
        npm install -g npm@latest
        
        # Install n module for Node.js version management
        npm install -g n
        
        # Install latest stable Node.js version using n
        n stable
        
        # Reload shell environment to use new Node.js version
        export PATH="$PATH"
        
        log_success "Node.js wurde erfolgreich installiert/aktualisiert."
        log_info "Installed Node.js version: $(node --version)"
        log_info "Installed npm version: $(npm --version)"
    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 \
            screen vim nano git \
            curl wget tree ncdu \
            gcc g++ build-essential
        
        # 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
}
# Function to display system overview
display_system_overview() {
    log_info "=== System Overview ==="
    echo -e "\n${BLUE}[Network Information]${NC}"
    echo "IP Addresses:"
    ip -4 addr show | grep inet | awk '{print $2}' | cut -d/ -f1

    echo -e "\n${BLUE}[Disk Usage]${NC}"
    df -h / | tail -n 1 | awk '{print "Root partition: " $4 " free of " $2}'

    echo -e "\n${BLUE}[Memory Usage]${NC}"
    free -h | awk '/^Mem:/ {print "RAM: " $3 " used of " $2 " total"}'

    echo -e "\n${BLUE}[CPU Information]${NC}"
    lscpu | grep "Model name" | cut -d: -f2- | sed 's/^[ \t]*//'
    echo "CPU Usage: $(top -bn1 | grep "Cpu(s)" | awk '{print $2}')%"

    echo -e "\n${BLUE}[Installed Components]${NC}"
    if command_exists docker; then
        echo "✓ Docker $(docker --version | cut -d' ' -f3 | tr -d ',')"
        if docker ps | grep -q dockge; then
            echo "✓ Dockge (running on port 5001)"
        fi
    fi
    if command_exists node; then
        echo "✓ Node.js $(node --version)"
    fi
    if command_exists ufw; then
        echo "✓ UFW ($(ufw status | grep Status | cut -d' ' -f2))"
    fi
    echo -e "\n${BLUE}[Open Ports]${NC}"
    netstat -tuln | grep LISTEN | awk '{print $4}' | sed 's/.*://' | tr '\n' ', ' | sed 's/,$/\n/'

    echo -e "\n${GREEN}Setup completed! System is ready to use.${NC}"
    echo -e "${YELLOW}Important: Please reconnect to your server via SSH to ensure all changes are properly applied.${NC}\n"
}

# Main execution
main() {
    setup_ssh
    update_system
    setup_basic_system
    install_docker
    install_nodejs
    display_system_overview
}

# Run main function
main