#!/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; }

# =============================
# CONFIG
# =============================
GIST_RAW_BASE="https://opengist.rmrf.online/weehong/f0d940c3c1214bf5b7996195199fdc09/raw/HEAD"

CONFIG_FILES=(
    ".alias"
    ".func"
    ".pathrc"
    ".sourcerc"
    ".vimrc"
    ".zshrc"
    ".config/starship.toml"
)

# =============================
# REQUIREMENTS
# =============================
check_requirements() {
    if [[ $EUID -eq 0 ]]; then
        err "Do not run as root on macOS"
        exit 1
    fi
}

# =============================
# INSTALLATION FUNCTIONS
# =============================
install_plugins() {
    log "Installing manual Zsh plugins..."
    local dir="$HOME/.zsh"
    mkdir -p "$dir"

    [[ -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 directly to $dir"
}

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

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
        # Strip the leading dot for the download URL
        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 plugins and Starship..."

    touch "$zshrc"

    # Add Starship init if not present
    if ! grep -q 'starship init zsh' "$zshrc"; then
        echo -e '\n# Initialize Starship' >> "$zshrc"
        echo 'eval "$(starship init zsh)"' >> "$zshrc"
    fi

    # Add autosuggestions if not present
    if ! grep -q 'zsh-autosuggestions.zsh' "$zshrc"; then
        echo -e '\n# Load Plugins' >> "$zshrc"
        echo 'source ~/.zsh/zsh-autosuggestions/zsh-autosuggestions.zsh' >> "$zshrc"
    fi

    # Add syntax highlighting if not present (MUST be last)
    if ! grep -q 'zsh-syntax-highlighting.zsh' "$zshrc"; then
        echo 'source ~/.zsh/zsh-syntax-highlighting/zsh-syntax-highlighting.zsh' >> "$zshrc"
    fi

    ok ".zshrc configured for zero-framework setup!"
    ok "Load order is preserved: .sourcerc -> .func -> .pathrc -> .alias"
}

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 "macOS Minimal Zsh Setup - Choose what to do"
    echo "==========================================="
    echo " 0) Run ALL steps (1-5)"
    echo " 1) Install manual Zsh plugins (autosuggestions, syntax highlighting)"
    echo " 2) Install Starship prompt"
    echo " 3) Download custom configs (~/.alias, .func, .vimrc, etc.)"
    echo " 4) Configure ~/.zshrc (Inject Starship & Plugins)"
    echo " 5) Switch to Zsh (Temporary Sub-shell)"
    echo " 6) 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)
        else
            to_run+=("$item")
        fi
    done

    if [[ ${#to_run[@]} -gt 0 ]]; then
        for choice in "${to_run[@]}"; do
            local skip=false

            if [[ ${#to_exclude[@]} -gt 0 ]]; then
                for ex in "${to_exclude[@]}"; do
                    if [[ "$choice" == "$ex" ]]; then
                        skip=true
                        break
                    fi
                done
            fi

            $skip && continue

            case "$choice" in
                1) install_plugins ;;
                2) install_starship ;;
                3) download_configs ;;
                4) configure_zshrc ;;
                5) switch_shell ;;
                6) exit 0 ;;
                *) warn "Skipping invalid option: $choice" ;;
            esac
            echo
        done
    fi
}

# =============================
# 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 "macOS minimal Zsh configuration complete!"
}

main "$@"