#!/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_SHELL_CHANGE=false
INSTALL_HOMEBREW=false

GIST_RAW_BASE="https://opengist.rmrf.online/weehong/f0d940c3c1214bf5b7996195199fdc09/raw/HEAD"
CONFIG_FILES=(
    ".alias"
    ".func"
    ".pathrc"
    ".sourcerc"
    ".vimrc"
    ".zshrc"
    ".config/starship.toml"
)

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

is_wsl() {
    grep -qi "microsoft" /proc/version 2>/dev/null
}

# =============================
# REQUIREMENTS
# =============================
check_requirements() {
    if [[ $EUID -eq 0 ]]; then
        err "Do not run as root. The script will request sudo when necessary."
        exit 1
    fi
    if ! command -v sudo >/dev/null 2>&1; then
        err "sudo required"
        exit 1
    fi
    if ! is_wsl; then
        err "This installer is intended for WSL"
        exit 1
    fi
}

# =============================
# INSTALLATION FUNCTIONS
# =============================
update_system() {
    $SKIP_PACKAGES && return
    log "Updating system..."
    local os=$(detect_os)
    
    case "$os" in
        ubuntu|debian)
            sudo apt-get update -y
            if ! sudo apt-get upgrade -y; then
                warn "System upgrade did not complete. This can happen when apt wants to downgrade a package."
                warn "Continuing because the Zsh setup does not require OS package upgrades to finish."
            fi
            ;;
        fedora)
            sudo dnf upgrade -y || warn "System upgrade failed; continuing installer"
            ;;
        arch)
            sudo pacman -Syu --noconfirm || warn "System upgrade failed; continuing installer"
            ;;
        *) warn "Auto-update not supported for OS: $os" ;;
    esac
    ok "System update step finished"
}

install_packages() {
    $SKIP_PACKAGES && return
    log "Installing core packages..."
    local os=$(detect_os)
    
    case "$os" in
        ubuntu|debian) sudo apt-get install -y zsh git vim curl wget unzip zip build-essential xz-utils ;;
        fedora) sudo dnf install -y zsh git vim curl wget unzip zip @development-tools xz ;;
        arch) sudo pacman -S --noconfirm zsh git vim curl wget unzip zip base-devel xz ;;
        *) warn "Auto-install not supported for OS: $os. Please install zsh, git, vim, curl manually." ;;
    esac
    ok "Packages installed"
}

set_timezone() {
    log "Checking timezone configuration..."
    if command -v timedatectl >/dev/null 2>&1 && timedatectl status >/dev/null 2>&1; then
        sudo timedatectl set-timezone Asia/Singapore
        ok "Timezone set to Asia/Singapore"
        return
    fi

    warn "Skipping timezone change: timedatectl is not available in this WSL environment"
}

install_homebrew() {
    ! $INSTALL_HOMEBREW && return
    command -v brew >/dev/null 2>&1 && 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)"
    
    if ! grep -qx "$zsh_path" /etc/shells; then
        echo "$zsh_path" | sudo tee -a /etc/shells >/dev/null
    fi
    chsh -s "$zsh_path"
    ok "Shell changed (open a new WSL session for it to take effect)"
}

install_oh_my_zsh() {
    log "Installing Oh My Zsh and plugins..."
    export RUNZSH=no
    export CHSH=no
    export KEEP_ZSHRC=yes

    if [[ ! -d "$HOME/.oh-my-zsh" ]]; then
        sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" --unattended
    else
        ok "Oh My Zsh already installed"
    fi

    local custom="${ZSH_CUSTOM:-$HOME/.oh-my-zsh/custom}"
    local plugin_dir="$custom/plugins"
    mkdir -p "$plugin_dir"

    [[ -d "$plugin_dir/zsh-autosuggestions" ]] || git clone https://github.com/zsh-users/zsh-autosuggestions "$plugin_dir/zsh-autosuggestions"
    [[ -d "$plugin_dir/zsh-syntax-highlighting" ]] || git clone https://github.com/zsh-users/zsh-syntax-highlighting "$plugin_dir/zsh-syntax-highlighting"

    ok "Oh My Zsh plugins installed to $plugin_dir"
}

install_starship() {
    if command -v starship >/dev/null 2>&1 || [[ -f "$HOME/.local/bin/starship" ]]; then
        ok "Starship already installed"
    else
        log "Installing Starship prompt..."
        if command -v brew >/dev/null 2>&1; then
            brew install starship
        else
            mkdir -p "$HOME/.local/bin"
            curl -fsSL https://starship.rs/install.sh | sh -s -- -y -b "$HOME/.local/bin"
        fi
        ok "Starship installed"
    fi
}

download_configs() {
    log "Downloading custom config files from OpenGist..."
    local backup="$HOME/.config_backup_$(date +%Y%m%d_%H%M%S)"
    mkdir -p "$backup"

    for f in "${CONFIG_FILES[@]}"; do
        local remote_file="${f#.}"
        [[ "$f" == ".config/starship.toml" ]] && remote_file="starship.toml"
        local url="$GIST_RAW_BASE/$remote_file"
        local target="$HOME/$f"
        local tmp="${target}.tmp.$$"

        mkdir -p "$(dirname "$target")"

        if [[ -f "$target" ]]; then
            cp "$target" "$backup/"
        fi

        log "Fetching $remote_file -> $f ..."
        if curl -fsSL "$url" -o "$tmp"; then
            mv "$tmp" "$target"
        else
            rm -f "$tmp"
            warn "Failed to download $remote_file"
        fi
    done

    ok "Configs downloaded (Backup at $backup)"
}

configure_zshrc() {
    local zshrc="$HOME/.zshrc"
    log "Configuring .zshrc for Oh My Zsh and Starship..."

    touch "$zshrc"
    cp "$zshrc" "$HOME/.zshrc.backup_$(date +%Y%m%d_%H%M%S)"

    cat > "$zshrc" <<'EOF'
# =============================================================================
# 1. HELPER FUNCTIONS & PATH
# =============================================================================
source_if_readable() {
    local file="$1"
    if [[ -f "$file" && -r "$file" ]]; then
        source "$file"
    fi
}

export PATH="$HOME/.local/bin:$PATH"

# =============================================================================
# 2. OH MY ZSH FRAMEWORK
# =============================================================================
export ZSH="${ZSH:-$HOME/.oh-my-zsh}"
ZSH_THEME=""

plugins=(
    git
    zsh-autosuggestions
    zsh-syntax-highlighting
)

source_if_readable "$ZSH/oh-my-zsh.sh"

# =============================================================================
# 3. THIRD-PARTY INITIALIZATION & CUSTOM CONFIGS
# =============================================================================
source_if_readable "$HOME/.sourcerc"
source_if_readable "$HOME/.func"
source_if_readable "$HOME/.pathrc"
source_if_readable "$HOME/.alias"

# =============================================================================
# 4. STARSHIP PROMPT
# =============================================================================
if [[ -z "${STARSHIP_CONFIG:-}" && -f "$HOME/.config/starship.toml" ]]; then
    export STARSHIP_CONFIG="$HOME/.config/starship.toml"
fi

if command -v starship >/dev/null 2>&1; then
    eval "$(starship init zsh)"
elif [[ -x "$HOME/.local/bin/starship" ]]; then
    eval "$("$HOME/.local/bin/starship" init zsh)"
fi
EOF

    ok ".zshrc configured for Oh My Zsh framework with Starship prompt"
}

switch_shell() {
    log "Starting Zsh session..."
    echo -e "${YELLOW}Type 'exit' to return to this installer menu.${NC}"
    echo "----------------------------------------"
    zsh -l
    echo "----------------------------------------"
    ok "Returned from Zsh session"
}

# =============================
# INTERACTIVE MENU
# =============================
show_menu() {
    echo "==========================================="
    echo "WSL Minimal Zsh Installer - Choose what to do"
    echo "==========================================="
    echo " 0) Run ALL steps (1-10)"
    echo " 1) Update system packages"
    echo " 2) Install core packages (zsh, git, vim, etc.)"
    echo " 3) Set Timezone (best effort)"
    echo " 4) Install Homebrew"
    echo " 5) Configure shell (chsh - sets default shell)"
    echo " 6) Install Oh My Zsh + plugins"
    echo " 7) Install Starship prompt"
    echo " 8) Download custom configs (from OpenGist)"
    echo " 9) Configure ~/.zshrc (Oh My Zsh + Starship)"
    echo "10) Switch to Zsh (Temporary Sub-shell)"
    echo "11) Quit"
    echo "==========================================="
}

run_choices() {
    local input
    read -p "Select: " input
    input="${input//,/ }"

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

    for item in $input; do
        if [[ "$item" == !* ]]; then
            to_exclude+=("${item:1}")
        elif [[ "$item" == "0" ]]; then
            to_run+=(1 2 3 4 5 6 7 8 9 10)
        else
            to_run+=("$item")
        fi
    done

    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_homebrew ;;
            5) configure_shell ;;
            6) install_oh_my_zsh ;;
            7) install_starship ;;
            8) download_configs ;;
            9) configure_zshrc ;;
            10) switch_shell ;;
            11) 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 "WSL Zsh installation/configuration complete!"
}

main "$@"