#!/bin/bash

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

# Configuration for UI
MAX_LOG_LINES=8
declare -a ACTIVE_STEPS=()

teardown_ui() {
    local ROWS
    ROWS=$(tput lines 2>/dev/null || echo 24)
    printf "\033[1;${ROWS}r"  # Reset scrolling region to full terminal
    printf "\033[?25h"        # Ensure cursor is visible
    printf "\033[${ROWS};1H"  # Move cursor to bottom explicitly
}
trap 'teardown_ui' EXIT INT TERM

setup_ui() {
    clear
    local ROWS
    ROWS=$(tput lines 2>/dev/null || echo 24)
    
    local num_steps=${#ACTIVE_STEPS[@]}
    local LOG_END=$ROWS
    local LOG_START=$(( LOG_END - MAX_LOG_LINES + 1 ))
    
    local track_end=$(( 3 + num_steps ))
    local MIN_LOG_START=$(( track_end + 3 ))
    
    if [ "$LOG_START" -le "$(( MIN_LOG_START - 1 ))" ]; then 
        LOG_START=$MIN_LOG_START
    fi
    
    SEP_LINE=$(( LOG_START - 1 ))
    INFO_LINE=$(( SEP_LINE - 1 ))
    
    printf "\033[1;1H${BLUE}=========================================${NC}\n"
    printf "\033[2;1H${CYAN}         HOMELOGIC SETUP SCRIPT          ${NC}\n"
    printf "\033[3;1H${BLUE}=========================================${NC}\n"
    
    for i in "${!ACTIVE_STEPS[@]}"; do
        local line=$(( 4 + i ))
        printf "\033[${line};1H  [          ] ${ACTIVE_STEPS[$i]}\n"
        update_step_status "${ACTIVE_STEPS[$i]}" "Pending"
    done
    
    printf "\033[${SEP_LINE};1H${YELLOW}--- Verbose Logs ------------------------${NC}\n"
    
    printf "\033[${LOG_START};${LOG_END}r" 
    printf "\033[${LOG_END};1H"
}

get_line_for_step() {
    for i in "${!ACTIVE_STEPS[@]}"; do
        if [ "${ACTIVE_STEPS[$i]}" = "$1" ]; then
            echo $(( 4 + i ))
            return
        fi
    done
    echo 0
}

update_step_status() {
    local step_name="$1"
    local status="$2"
    local line=$(get_line_for_step "$step_name")
    
    if [ "$line" -gt 0 ]; then
        local color="${NC}"
        case "$status" in
            Pending) color="${NC}" ;;
            Running) color="${BLUE}" ;;
            Installed) color="${GREEN}" ;;
            Skipped|Removed) color="${YELLOW}" ;;
        esac
        printf "\0337\033[%d;4H${color}%-10s${NC}\0338" "$line" "$status"
    fi
}

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

print_info() {
    local IL=${INFO_LINE:-10}
    printf "\0337\033[${IL};1H\033[K${CYAN}  ➜ %s${NC}\0338" "$1"
    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."
        update_step_status "$step_name" "Skipped"
        ;;
      * )
        update_step_status "$step_name" "Running"
        $callback
        ;;
    esac
}

setup_server() {
    update_step_status "Server" "Running"
    print_step "Setting up Server"
    print_info "Updating packages and configuring firewall..."
    sudo apt-get update
    sudo ufw allow 22/tcp #ssh
    sudo ufw allow 80/tcp #nginx
    sudo ufw allow 81/tcp #nginx
    sudo ufw allow 443/tcp #nginx
    sudo ufw allow 8123/tcp #homeassistant
    sudo ufw allow 8123/udp #homeassistant
    sudo ufw allow 3000/tcp #nextcloud
    sudo ufw allow 3001/tcp #zigbee2mqtt
    sudo ufw allow 1883/tcp #mosquitto
    sudo ufw allow 2001/tcp #arcane
    update_step_status "Server" "Installed"
}

setup_docker() {
    update_step_status "Docker" "Running"
    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
    update_step_status "Docker" "Installed"
}

setup_zsh() {
    update_step_status "ZSH" "Running"
    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)"
    update_step_status "ZSH" "Installed"
}

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
    update_step_status "Tailscale" "Installed"
}

setup_wifitui() {
    print_step "Setting up Wifitui"
    print_info "Installing Network Manager..."
    sudo apt-get install network-manager -y
    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 -y /tmp/wifitui.deb
    rm -f /tmp/wifitui.deb
    update_step_status "Wifitui" "Installed"
}

setup_ha_watchdog() {
    local watchdog_path="$(cd "$(dirname "$0")" && pwd)/ha_watchdog.sh"
    
    print_step "Setting up HA Watchdog"
    if [ -f "/etc/cron.d/ha_watchdog" ]; then
        print_info "Uninstalling HA Watchdog..."
        sudo rm -f "/etc/cron.d/ha_watchdog"
        sudo rm -f "/usr/local/bin/ha_watchdog.sh" 2>/dev/null || true
        update_step_status "HA Watchdog" "Removed"
    else
        print_info "Configuring HA Watchdog script permissions..."
        chmod +x "$watchdog_path"
        
        print_info "Creating cron job (every 5 minutes)..."
        echo "*/5 * * * * root $watchdog_path" | sudo tee /etc/cron.d/ha_watchdog > /dev/null
        
        update_step_status "HA Watchdog" "Installed"
    fi
}

main() {
    clear
    printf "\n${BLUE}=========================================${NC}\n"
    printf "${CYAN}         HOMELOGIC SETUP MENU            ${NC}\n"
    printf "${BLUE}=========================================${NC}\n"
    printf "  1) Run ALL Steps (Recommended)\n"
    printf "  2) Run Main Steps Only\n"
    printf "  3) Run Optional Steps Only\n"
    printf "  4) Install/Uninstall HA Watchdog\n"
    printf "  q) Exit\n"
    printf "${BLUE}=========================================${NC}\n"
    printf "${YELLOW}Select an option [1-4 or q]: ${NC}"
    
    read opt
    case "$opt" in
        exit|quit|q|Q|0)
            printf "\nExiting setup.\n"
            exit 0
            ;;
        1)
            ACTIVE_STEPS=("Server" "Docker" "ZSH" "Tailscale" "Wifitui")
            setup_ui
            setup_server
            setup_docker
            setup_zsh
            ask_optional_step "Tailscale" setup_tailscale
            ask_optional_step "Wifitui" setup_wifitui
            ;;
        2) 
            ACTIVE_STEPS=("Server" "Docker" "ZSH")
            setup_ui
            setup_server
            setup_docker
            setup_zsh
            ;;
        3) 
            ACTIVE_STEPS=("Tailscale" "Wifitui")
            setup_ui
            ask_optional_step "Tailscale" setup_tailscale
            ask_optional_step "Wifitui" setup_wifitui
            ;;
        4)
            ACTIVE_STEPS=("HA Watchdog")
            setup_ui
            setup_ha_watchdog
            ;;
        *)
            printf "\n${YELLOW}⚠ Invalid option selected.${NC}\n\n"
            sleep 1
            main
            return
            ;;
    esac
    
    teardown_ui
    printf "\n${GREEN}✔ Task complete! Please log out and log back in for shell changes to take effect.${NC}\n"
}

# Execute main process
main "$@"