#!/usr/bin/env bash
#══════════════════════════════════════════════════════════════
#  Zsh + Oh-My-Zsh Installer — Ubuntu 24.04
#  Plugins: git zsh-autosuggestions zsh-syntax-highlighting
#  Self-extracting offline installer (with online fallback)
#══════════════════════════════════════════════════════════════

set -e

RST="\033[0m"; BOLD="\033[1m"; DIM="\033[2m"
GREEN="\033[32m"; RED="\033[31m"; YELLOW="\033[33m"; CYAN="\033[36m"

TARGET_USER="${SUDO_USER:-$USER}"
TARGET_HOME="$(getent passwd "$TARGET_USER" | cut -d: -f6)"
ZSH_DIR="$TARGET_HOME/.oh-my-zsh"
CUSTOM_DIR="$ZSH_DIR/custom"
PLUGIN_DIR="$CUSTOM_DIR/plugins"
ZSHRC="$TARGET_HOME/.zshrc"

log()  { echo -e "  ${CYAN}▸${RST} $*"; }
ok()   { echo -e "  ${GREEN}✔${RST} $*"; }
warn() { echo -e "  ${YELLOW}!${RST} $*"; }
err()  { echo -e "  ${RED}✖${RST} $*" >&2; }

need_root() {
    if [[ $EUID -ne 0 ]]; then
        err "Please run with sudo: sudo bash $0"
        exit 1
    fi
}

# ── Self-extracting payload ──────────────────────────────────
_self_path()   { readlink -f "$0" 2>/dev/null || echo "$0"; }
_has_payload() { grep -q '^__PAYLOAD__$' "$(_self_path)" 2>/dev/null; }

_extract_payload() {
    local self dest marker
    self="$(_self_path)"
    dest="$1"
    marker="$(grep -an '^__PAYLOAD__$' "$self" | tail -1 | cut -d: -f1)"
    [[ -z "$marker" ]] && return 1
    tail -n +"$(( marker + 1 ))" "$self" | tar xz -C "$dest"
}

# ── Install zsh ──────────────────────────────────────────────
install_zsh() {
    if command -v zsh >/dev/null 2>&1; then
        ok "zsh already installed: $(zsh --version)"
        return 0
    fi

    if [[ -d "$1/deb" ]] && ls "$1"/deb/*.deb >/dev/null 2>&1; then
        log "Installing zsh from bundled .deb packages..."
        DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends "$1"/deb/*.deb \
            || dpkg -i "$1"/deb/*.deb || { apt-get -f install -y; dpkg -i "$1"/deb/*.deb; }
    else
        log "Installing zsh via apt..."
        apt-get update -qq
        DEBIAN_FRONTEND=noninteractive apt-get install -y zsh git curl ca-certificates
    fi
    ok "zsh installed: $(zsh --version)"
}

# ── Install oh-my-zsh ────────────────────────────────────────
install_omz() {
    local src="$1"
    if [[ -d "$ZSH_DIR" ]]; then
        ok "oh-my-zsh already at $ZSH_DIR"
        return 0
    fi
    log "Installing oh-my-zsh to $ZSH_DIR ..."
    if [[ -d "$src/oh-my-zsh" ]]; then
        cp -a "$src/oh-my-zsh" "$ZSH_DIR"
    else
        warn "Bundled oh-my-zsh not found, cloning from GitHub..."
        sudo -u "$TARGET_USER" git clone --depth=1 \
            https://github.com/ohmyzsh/ohmyzsh.git "$ZSH_DIR"
    fi
    chown -R "$TARGET_USER:$TARGET_USER" "$ZSH_DIR"
    ok "oh-my-zsh ready"
}

# ── Install plugins ──────────────────────────────────────────
install_plugin() {
    local name="$1" src_dir="$2" repo="$3"
    local dst="$PLUGIN_DIR/$name"
    if [[ -d "$dst" ]]; then
        ok "plugin already installed: $name"
        return 0
    fi
    mkdir -p "$PLUGIN_DIR"
    if [[ -d "$src_dir/$name" ]]; then
        log "Installing plugin: $name (bundled)"
        cp -a "$src_dir/$name" "$dst"
    else
        log "Installing plugin: $name (clone)"
        sudo -u "$TARGET_USER" git clone --depth=1 "$repo" "$dst"
    fi
    chown -R "$TARGET_USER:$TARGET_USER" "$dst"
}

# ── Configure .zshrc ─────────────────────────────────────────
configure_zshrc() {
    log "Configuring $ZSHRC ..."
    if [[ ! -f "$ZSHRC" ]]; then
        if [[ -f "$ZSH_DIR/templates/zshrc.zsh-template" ]]; then
            cp "$ZSH_DIR/templates/zshrc.zsh-template" "$ZSHRC"
        else
            : > "$ZSHRC"
        fi
    else
        cp "$ZSHRC" "$ZSHRC.bak.$(date +%s)"
        ok "Existing .zshrc backed up"
    fi

    # Ensure ZSH points at our install
    if grep -q '^export ZSH=' "$ZSHRC"; then
        sed -i "s|^export ZSH=.*|export ZSH=\"$ZSH_DIR\"|" "$ZSHRC"
    else
        sed -i "1i export ZSH=\"$ZSH_DIR\"" "$ZSHRC"
    fi

    # Set plugins line
    local plugins_line='plugins=(git zsh-autosuggestions zsh-syntax-highlighting)'
    if grep -qE '^plugins=\(' "$ZSHRC"; then
        sed -i "s|^plugins=(.*)|$plugins_line|" "$ZSHRC"
    else
        echo "$plugins_line" >> "$ZSHRC"
    fi

    # Make sure oh-my-zsh.sh is sourced
    if ! grep -q 'source \$ZSH/oh-my-zsh.sh' "$ZSHRC" \
       && ! grep -q 'source $ZSH/oh-my-zsh.sh' "$ZSHRC"; then
        echo 'source $ZSH/oh-my-zsh.sh' >> "$ZSHRC"
    fi

    chown "$TARGET_USER:$TARGET_USER" "$ZSHRC"
    ok ".zshrc configured"
}

# ── Set default shell ────────────────────────────────────────
set_default_shell() {
    local zsh_bin
    zsh_bin="$(command -v zsh)"
    local cur
    cur="$(getent passwd "$TARGET_USER" | cut -d: -f7)"
    if [[ "$cur" == "$zsh_bin" ]]; then
        ok "Default shell is already zsh"
        return 0
    fi
    log "Changing default shell for $TARGET_USER to $zsh_bin ..."
    chsh -s "$zsh_bin" "$TARGET_USER"
    ok "Default shell set (re-login to apply)"
}

# ── Main ─────────────────────────────────────────────────────
main() {
    need_root
    echo -e "${BOLD}${CYAN}══════════════════════════════════════════${RST}"
    echo -e "${BOLD}  Zsh Installer — Ubuntu 24.04${RST}"
    echo -e "${BOLD}${CYAN}══════════════════════════════════════════${RST}"
    echo ""
    log "Target user : ${BOLD}$TARGET_USER${RST}"
    log "Target home : ${BOLD}$TARGET_HOME${RST}"
    echo ""

    local work=""
    if _has_payload; then
        work="$(mktemp -d)"
        trap 'rm -rf "$work"' EXIT
        log "Extracting embedded payload..."
        _extract_payload "$work" || { err "Payload extraction failed"; exit 1; }
        ok "Payload extracted to $work"
    else
        warn "No embedded payload — will download from network"
        apt-get update -qq
        DEBIAN_FRONTEND=noninteractive apt-get install -y git curl ca-certificates
    fi

    install_zsh "$work"
    install_omz "$work"
    install_plugin "zsh-autosuggestions" "$work" \
        "https://github.com/zsh-users/zsh-autosuggestions"
    install_plugin "zsh-syntax-highlighting" "$work" \
        "https://github.com/zsh-users/zsh-syntax-highlighting"
    configure_zshrc
    set_default_shell

    echo ""
    echo -e "  ${GREEN}${BOLD}✔ Done.${RST}"
    echo -e "  ${DIM}Run 'zsh' to start now, or log out and back in.${RST}"
    echo ""
}

main "$@"
exit 0