#!/bin/bash
set -euo pipefail

# =============================
# COLORS & LOGGING
# =============================
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log()  { echo -e "${BLUE}[INFO]${NC} $*"; }
ok()   { echo -e "${GREEN}[OK]${NC} $*"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
err()  { echo -e "${RED}[ERROR]${NC} $*" >&2; }

# =============================
# FLAGS & CONFIG
# =============================
SKIP_PACKAGES=false
SKIP_GIT_CONFIG=false
SKIP_SHELL_CHANGE=false
CUSTOM_CONFIG=false 
SKIP_XCLIP=false
INSTALL_HOMEBREW=false

declare -A CONFIG_FILES=(
    [".alias"]="https://gist.githubusercontent.com/weehongkoh/72bdb76beacacf2ca3dd39a72395b9ee/raw/alias"
    [".func"]="https://gist.githubusercontent.com/weehongkoh/72bdb76beacacf2ca3dd39a72395b9ee/raw/func"
    [".pathrc"]="https://gist.githubusercontent.com/weehongkoh/72bdb76beacacf2ca3dd39a72395b9ee/raw/pathrc"
    [".sourcerc"]="https://gist.githubusercontent.com/weehongkoh/72bdb76beacacf2ca3dd39a72395b9ee/raw/sourcerc"
    [".vimrc"]="https://gist.githubusercontent.com/weehongkoh/72bdb76beacacf2ca3dd39a72395b9ee/raw/vimrc"
    [".zshrc"]="https://gist.githubusercontent.com/weehongkoh/72bdb76beacacf2ca3dd39a72395b9ee/raw/zshrc"
)

# =============================
# OS DETECTION
# =============================
detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        echo "$ID"
    else
        echo "unknown"
    fi
}

# =============================
# REQUIREMENTS
# =============================
check_requirements() {
    # Fixed to prevent silent failure with 'set -e'
    if [[ $EUID -eq 0 ]]; then
        err "Do not run as root"
        exit 1
    fi
    if ! command -v sudo >/dev/null; then
        err "sudo required"
        exit 1
    fi
}

# =============================
# INSTALLATION FUNCTIONS
# =============================
update_system() {
    $SKIP_PACKAGES && return
    log "Updating system..."
    sudo apt-get update -y
    sudo apt-get upgrade -y
    ok "System updated"
}

install_packages() {
    $SKIP_PACKAGES && return
    log "Installing core packages..."
    sudo apt-get install -y \
        zsh git vim curl wget unzip zip build-essential xz-utils
    ok "Packages installed"
}

set_timezone() {
    log "Setting timezone to Asia/Singapore..."
    sudo timedatectl set-timezone Asia/Singapore
    ok "Timezone set to Asia/Singapore"
}

install_xclip() {
    local os
    os=$(detect_os)
    [[ "$os" != "ubuntu" && "$os" != "debian" ]] && return 0
    $SKIP_XCLIP && return 0
    command -v xclip >/dev/null && ok "xclip already installed" && return 0

    read -r -p "Install xclip? (y/N): " r
    case "$r" in
        [Yy]) sudo apt-get install -y xclip; ok "xclip installed" ;;
        *) log "Skipping xclip installation" ;;
    esac
}

configure_git() {
    $SKIP_GIT_CONFIG && return
    log "Configuring Git..."
    read -p "Git name (leave empty to skip): " name
    read -p "Git email (leave empty to skip): " email
    [[ -n "$name" ]] && git config --global user.name "$name"
    [[ -n "$email" ]] && git config --global user.email "$email"
    ok "Git configured"
}

install_homebrew() {
    ! $INSTALL_HOMEBREW && return
    command -v brew >/dev/null && ok "Homebrew already installed" && return
    log "Installing Homebrew..."
    /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
    ok "Homebrew installed"
}

configure_shell() {
    $SKIP_SHELL_CHANGE && return
    log "Changing default shell to zsh..."
    local zsh_path
    zsh_path=$(command -v zsh)
    grep -qx "$zsh_path" /etc/shells || echo "$zsh_path" | sudo tee -a /etc/shells >/dev/null
    sudo chsh -s "$zsh_path" "$USER"
    ok "Shell changed (requires logout/login to take effect)"
}

install_oh_my_zsh() {
    [[ -d "$HOME/.oh-my-zsh" ]] && ok "Oh My Zsh already installed" && return
    log "Installing Oh My Zsh..."
    # Keep zshrc ensures we don't blow away configs if they exist
    RUNZSH=no CHSH=no KEEP_ZSHRC=yes \
        sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
    ok "Oh My Zsh installed"
}

install_plugins() {
    log "Installing plugins..."
    local dir="${ZSH_CUSTOM:-$HOME/.oh-my-zsh/custom}/plugins"
    [[ -d "$dir/zsh-autosuggestions" ]] || git clone https://github.com/zsh-users/zsh-autosuggestions "$dir/zsh-autosuggestions"
    [[ -d "$dir/zsh-syntax-highlighting" ]] || git clone https://github.com/zsh-users/zsh-syntax-highlighting "$dir/zsh-syntax-highlighting"
    ok "Plugins installed"
}

install_theme() {
    log "Installing Spaceship theme..."
    local themes="${ZSH_CUSTOM:-$HOME/.oh-my-zsh/custom}/themes"
    local dir="$themes/spaceship-prompt"
    [[ -d "$dir" ]] || git clone https://github.com/spaceship-prompt/spaceship-prompt.git "$dir" --depth=1
    ln -sf "$dir/spaceship.zsh-theme" "$themes/spaceship.zsh-theme"
    ok "Theme installed"
}

download_configs() {
    log "Downloading custom config files..."
    local backup="$HOME/.config_backup_$(date +%Y%m%d_%H%M%S)"
    mkdir -p "$backup"
    for f in "${!CONFIG_FILES[@]}"; do
        if [[ -f "$HOME/$f" ]]; then
             cp "$HOME/$f" "$backup/"
        fi
        log "Fetching $f ..."
        curl -fsSL "${CONFIG_FILES[$f]}" -o "$HOME/$f" || warn "Failed to download $f"
    done
    ok "Configs downloaded (Backup at $backup)"
}

update_zshrc() {
    local zshrc="$HOME/.zshrc"
    log "Updating .zshrc..."
    
    # Create zshrc if missing to prevent errors
    if [[ ! -f "$zshrc" ]]; then
        warn ".zshrc not found, creating new one..."
        touch "$zshrc"
    fi

    # Use || true to prevent 'set -e' from exiting if grep finds nothing
    grep -q 'ZSH_THEME="spaceship"' "$zshrc" || sed -i 's/ZSH_THEME=".*"/ZSH_THEME="spaceship"/' "$zshrc"
    grep -q 'zsh-autosuggestions' "$zshrc" || sed -i 's/plugins=(/plugins=(zsh-autosuggestions /' "$zshrc"
    grep -q 'zsh-syntax-highlighting' "$zshrc" || sed -i 's/plugins=(/plugins=(zsh-syntax-highlighting /' "$zshrc"
    ok ".zshrc updated with plugins and theme"
}

switch_shell() {
    $SKIP_SHELL_CHANGE && return
    
    log "Starting Zsh session..."
    echo -e "${YELLOW}Type 'exit' to return to this installer menu.${NC}"
    echo "----------------------------------------"
    
    # Run zsh as a subprocess, not exec
    zsh -l
    
    echo "----------------------------------------"
    ok "Returned from Zsh session"
}

# =============================
# INTERACTIVE MENU
# =============================
show_menu() {
    echo "==========================================="
    echo "Zsh Installer - Choose what to do"
    echo "==========================================="
    echo " 0) Run ALL steps (1-13)"
    echo " 1) Update system packages"
    echo " 2) Install core packages (zsh, git, vim, etc.)"
    echo " 3) Set Timezone (Asia/Singapore)"
    echo " 4) Install xclip"
    echo " 5) Configure Git"
    echo " 6) Install Homebrew"
    echo " 7) Configure shell (chsh - sets default shell)"
    echo " 8) Install Oh My Zsh"
    echo " 9) Install plugins (autosuggestions, syntax highlighting)"
    echo "10) Install Spaceship theme"
    echo "11) Download custom configs (~/.alias, .vimrc, etc.)"
    echo "12) Update ~/.zshrc for plugins & theme"
    echo "13) Switch to Zsh (Temporary Sub-shell)"
    echo "14) Quit"
    echo "==========================================="
    echo "Inputs: '0' (All), '2 3 7' (Specific), '0 !4 !6' (All except 4 and 6)"
}

run_choices() {
    local input
    read -p "Select: " input
    input="${input//,/ }" # replace commas with spaces

    local -a to_run=()
    local -a to_exclude=()

    # Parse positive selections and negative exclusions
    for item in $input; do
        if [[ "$item" == !* ]]; then
            to_exclude+=("${item:1}") # Strip the '!' character
        elif [[ "$item" == "0" ]]; then
            to_run+=(1 2 3 4 5 6 7 8 9 10 11 12 13)
        else
            to_run+=("$item")
        fi
    done

    # Loop through intended runs and apply exclusions
    for choice in "${to_run[@]}"; do
        
        local skip=false
        for ex in "${to_exclude[@]}"; do
            if [[ "$choice" == "$ex" ]]; then
                skip=true
                break
            fi
        done
        
        $skip && continue

        case "$choice" in
            1) update_system ;;
            2) install_packages ;;
            3) set_timezone ;;
            4) install_xclip ;;
            5) configure_git ;;
            6) install_homebrew ;;
            7) configure_shell ;;
            8) install_oh_my_zsh ;;
            9) install_plugins ;;
            10) install_theme ;;
            11) download_configs ;;
            12) update_zshrc ;;
            13) switch_shell ;;
            14) log "Exiting..."; exit 0 ;;
            *) warn "Skipping invalid option: $choice" ;;
        esac
        echo
    done
}

# =============================
# MAIN
# =============================
main() {
    check_requirements
    while true; do
        show_menu
        run_choices
        read -p "Do you want to run more options? (y/n): " again
        [[ "$again" =~ ^[Yy]$ ]] || break
    done
    ok "Zsh installation/configuration complete!"
}

main "$@"