#!/bin/bash

# Configuration for Colors
BLUE='\033[1;34m'
GREEN='\033[1;32m'
YELLOW='\033[1;33m'
CYAN='\033[1;36m'
NC='\033[0m'

print_step() {
    printf "\n${BLUE}========== %s ==========${NC}\n" "$1"
}
print_info() {
    printf "${CYAN}➜ %s${NC}\n" "$1"
}

ask_optional_step() {
    local step_name="$1"
    local callback="$2"
    
    printf "\n${YELLOW}❓ Do you want to install and setup %s? [Y/n] ${NC}" "$step_name"
    read choice
    case "$choice" in 
      n|N|no|No|NO )
        print_info "Skipping $step_name setup."
        ;;
      * )
        $callback
        ;;
    esac
}

setup_server() {
    print_step "Setting up Server"
    print_info "Updating packages and configuring firewall..."
    sudo apt-get update
    sudo ufw allow 22/tcp
    sudo ufw allow 443/tcp
    sudo ufw allow 8123/tcp
    sudo ufw allow 8123/udp
    sudo ufw allow 3000/tcp
    sudo ufw allow 3001/tcp
}

setup_docker() {
    print_step "Setting up Docker"
    print_info "Installing prerequisites..."
    sudo apt -y install lsb-release gnupg apt-transport-https ca-certificates curl software-properties-common
    if [ ! -f /etc/apt/trusted.gpg.d/docker.gpg ]; then
        print_info "Downloading Docker GPG key..."
        curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/trusted.gpg.d/docker.gpg
    fi
    print_info "Adding Docker repository..."
    sudo add-apt-repository -y "deb [arch=$(dpkg --print-architecture)] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable"
    sudo apt update
    print_info "Installing Docker CE..."
    sudo apt install -y docker-ce docker-ce-cli containerd.io docker-compose-plugin

    print_info "Configuring Docker service and permissions..."
    sudo systemctl enable --now docker.service
    if ! getent group docker > /dev/null 2>&1; then
        sudo groupadd docker
    fi
    if ! id -nG "$USER" | grep -qw "docker"; then
        sudo usermod -aG docker "$USER"
    fi
    sudo docker network create homelogic 2>/dev/null || true
}

setup_zsh() {
    print_step "Setting up ZSH"
    print_info "Configuring .zshrc and hushlogin..."
    touch ../.hushlogin
    cp -f .zshrc.base ../.zshrc

    sudo apt install -y zsh

    ZSH_PATH=$(command -v zsh)
    if [ "${SHELL##*/}" != "zsh" ] && [ -n "$ZSH_PATH" ]; then
        print_info "Changing default shell to ZSH..."
        chsh -s "$ZSH_PATH"
    fi
    print_info "Installing Oh My Zsh and plugins..."
    git clone https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions 2>/dev/null || true
    RUNZSH=no CHSH=no sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
}

setup_tailscale() {
    print_step "Setting up Tailscale"
    if ! command -v tailscale > /dev/null 2>&1; then
        print_info "Installing Tailscale..."
        curl -fsSL https://tailscale.com/install.sh | sh
        sudo tailscale up
    else
        print_info "Tailscale is already installed."
    fi
}

setup_wifitui() {
    print_step "Setting up Wifitui"
    print_info "Downloading Wifitui..."
    TAG=$(curl -s https://api.github.com/repos/shazow/wifitui/releases/latest | grep "tag_name" | cut -d '"' -f4)
    OS="linux-$(uname -m)" # x86_64 or arm64
    LATEST_RELEASE="https://github.com/shazow/wifitui/releases/download/${TAG}/wifitui-${TAG:1}-${OS}"
    curl -Lfs "${LATEST_RELEASE}.deb" -o /tmp/wifitui.deb
    sudo apt install /tmp/wifitui.deb
    rm -f /tmp/wifitui.deb
}

main() {
    setup_server
    setup_docker
    setup_zsh
    
    ask_optional_step "Tailscale" setup_tailscale
    ask_optional_step "Wifitui" setup_wifitui
    
    printf "\n${GREEN}✔ Setup complete! Please log out and log back in for all changes to take effect.${NC}\n"
}

# Execute main process
main "$@"