#!/usr/bin/env bash

#########################################################################################
#### chezmoi-bootstrap.sh — Bootstrap a new server with chezmoi-managed dotfiles      ####
#### Installs chezmoi, clones your dotfile repo, and applies configs in one step       ####
#### Dry-run by default — nothing is applied without --force                           ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./chezmoi-bootstrap.sh --repo https://github.com/user/dotfiles.git            ####
####   ./chezmoi-bootstrap.sh --repo git@github.com:user/dotfiles.git --force         ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
REPO_URL=""
INSTALL_DIR="/usr/local/bin"
DRY_RUN="${DRY_RUN:-true}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"
INSTALL_AGE="${INSTALL_AGE:-false}"
AGE_KEY_PATH="${AGE_KEY_PATH:-$HOME/.config/chezmoi/key.txt}"
PRE_PACKAGES=""
CHEZMOI_ARGS=""

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
ERRORS=0

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]] || [[ "$COLOR" == "auto" && ! -t 1 ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" RESET=""
    else
        RED='\033[0;31m'    GREEN='\033[0;32m'  YELLOW='\033[0;33m'
        BLUE='\033[0;34m'   CYAN='\033[0;36m'   BOLD='\033[1m'
        RESET='\033[0m'
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()   { echo -e "${GREEN}[INFO]${RESET}  $*"; }
warn()  { echo -e "${YELLOW}[WARN]${RESET}  $*" >&2; }
err()   { echo -e "${RED}[ERROR]${RESET} $*" >&2; ((ERRORS++)) || true; }
debug() { [[ "$VERBOSE" == "true" ]] && echo -e "${CYAN}[DEBUG]${RESET} $*"; }
step()  { echo -e "\n${BOLD}${BLUE}── $* ──${RESET}"; }

# ── Usage ─────────────────────────────────────────────────────────────
usage() {
    cat << EOF
${BOLD}$SCRIPT_NAME${RESET} — Bootstrap a new server with chezmoi dotfiles

${BOLD}USAGE${RESET}
    $SCRIPT_NAME --repo <url> [OPTIONS]

${BOLD}REQUIRED${RESET}
    --repo <url>            Git repository URL (HTTPS or SSH)

${BOLD}OPTIONS${RESET}
    --force                 Apply changes (default: dry-run)
    --install-dir <path>    Chezmoi install directory (default: /usr/local/bin)
    --install-age           Also install age for encrypted files
    --age-key <path>        Path to age key file (default: ~/.config/chezmoi/key.txt)
    --packages <list>       Comma-separated packages to install first
    --chezmoi-args <args>   Extra arguments to pass to chezmoi init
    --verbose               Show debug output
    --no-color              Disable colored output
    --help                  Show this help

${BOLD}EXAMPLES${RESET}
    # Dry run — see what would happen
    $SCRIPT_NAME --repo https://github.com/user/dotfiles.git

    # Apply dotfiles from a private repo
    $SCRIPT_NAME --repo git@github.com:user/dotfiles.git --force

    # Install age + pre-install packages + apply
    $SCRIPT_NAME --repo git@github.com:user/dotfiles.git \\
        --install-age --packages vim,tmux,htop --force

    # Custom install dir + verbose
    $SCRIPT_NAME --repo https://github.com/user/dotfiles.git \\
        --install-dir \$HOME/.local/bin --verbose --force

${BOLD}ENVIRONMENT VARIABLES${RESET}
    DRY_RUN         Set to false to apply (same as --force)
    VERBOSE         Set to true for debug output
    COLOR           auto | always | never
    AGE_KEY_PATH    Path to age identity file
EOF
}

# ── Argument Parsing ──────────────────────────────────────────────────
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --repo)         REPO_URL="$2"; shift 2 ;;
            --force)        DRY_RUN="false"; shift ;;
            --install-dir)  INSTALL_DIR="$2"; shift 2 ;;
            --install-age)  INSTALL_AGE="true"; shift ;;
            --age-key)      AGE_KEY_PATH="$2"; shift 2 ;;
            --packages)     PRE_PACKAGES="$2"; shift 2 ;;
            --chezmoi-args) CHEZMOI_ARGS="$2"; shift 2 ;;
            --verbose)      VERBOSE="true"; shift ;;
            --no-color)     COLOR="never"; shift ;;
            --help|-h)      usage; exit 0 ;;
            *)              err "Unknown option: $1"; usage; exit 1 ;;
        esac
    done

    if [[ -z "$REPO_URL" ]]; then
        err "Missing required --repo argument"
        usage
        exit 1
    fi
}

# ── System Detection ──────────────────────────────────────────────────
detect_system() {
    step "Detecting system"

    OS="$(uname -s)"
    ARCH="$(uname -m)"
    HOSTNAME_SHORT="$(hostname -s)"

    if [[ -f /etc/os-release ]]; then
        # shellcheck disable=SC1091
        . /etc/os-release
        DISTRO="${ID:-unknown}"
        DISTRO_VERSION="${VERSION_ID:-unknown}"
    else
        DISTRO="unknown"
        DISTRO_VERSION="unknown"
    fi

    # Detect package manager
    if command -v apt > /dev/null 2>&1; then
        PKG_MGR="apt"
    elif command -v dnf > /dev/null 2>&1; then
        PKG_MGR="dnf"
    elif command -v yum > /dev/null 2>&1; then
        PKG_MGR="yum"
    elif command -v pacman > /dev/null 2>&1; then
        PKG_MGR="pacman"
    else
        PKG_MGR="unknown"
    fi

    log "OS: $OS ($ARCH)"
    log "Distro: $DISTRO $DISTRO_VERSION"
    log "Package manager: $PKG_MGR"
    log "Hostname: $HOSTNAME_SHORT"
}

# ── Package Installation ─────────────────────────────────────────────
install_packages() {
    local packages="$1"
    if [[ -z "$packages" ]]; then
        return 0
    fi

    step "Installing prerequisite packages"

    # Convert comma-separated to space-separated
    local pkg_list
    pkg_list="${packages//,/ }"

    log "Packages: $pkg_list"

    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY RUN] Would install: $pkg_list"
        return 0
    fi

    case "$PKG_MGR" in
        apt)
            sudo apt update -qq
            # shellcheck disable=SC2086
            sudo apt install -y -qq $pkg_list
            ;;
        dnf|yum)
            # shellcheck disable=SC2086
            sudo "$PKG_MGR" install -y -q $pkg_list
            ;;
        pacman)
            # shellcheck disable=SC2086
            sudo pacman -S --noconfirm $pkg_list
            ;;
        *)
            warn "Unknown package manager — install manually: $pkg_list"
            ;;
    esac

    log "Packages installed"
}

# ── Install chezmoi ───────────────────────────────────────────────────
install_chezmoi() {
    step "Installing chezmoi"

    if command -v chezmoi > /dev/null 2>&1; then
        local current_version
        current_version="$(chezmoi --version | awk '{print $3}')"
        log "chezmoi already installed: $current_version"
        return 0
    fi

    log "Install directory: $INSTALL_DIR"

    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY RUN] Would install chezmoi to $INSTALL_DIR"
        return 0
    fi

    mkdir -p "$INSTALL_DIR"
    sh -c "$(curl -fsLS get.chezmoi.io)" -- -b "$INSTALL_DIR"

    if command -v chezmoi > /dev/null 2>&1; then
        log "chezmoi installed: $(chezmoi --version | awk '{print $3}')"
    else
        # Might not be in PATH yet
        if [[ -x "$INSTALL_DIR/chezmoi" ]]; then
            export PATH="$INSTALL_DIR:$PATH"
            log "chezmoi installed: $(chezmoi --version | awk '{print $3}')"
            log "Added $INSTALL_DIR to PATH"
        else
            err "chezmoi installation failed"
            return 1
        fi
    fi
}

# ── Install age ───────────────────────────────────────────────────────
install_age() {
    if [[ "$INSTALL_AGE" != "true" ]]; then
        return 0
    fi

    step "Installing age"

    if command -v age > /dev/null 2>&1; then
        log "age already installed: $(age --version)"
        return 0
    fi

    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY RUN] Would install age"
        return 0
    fi

    case "$PKG_MGR" in
        apt)    sudo apt install -y -qq age ;;
        dnf)    sudo dnf install -y -q age ;;
        *)
            # Fallback: install from GitHub
            local age_version
            age_version="$(curl -s https://api.github.com/repos/FiloSottile/age/releases/latest | grep tag_name | cut -d'"' -f4)"
            curl -fsSL "https://github.com/FiloSottile/age/releases/download/${age_version}/age-${age_version}-linux-amd64.tar.gz" | \
                sudo tar -xz -C /usr/local/bin/ --strip-components=1 age/age age/age-keygen
            ;;
    esac

    log "age installed: $(age --version)"
}

# ── Age Key Setup ─────────────────────────────────────────────────────
setup_age_key() {
    if [[ "$INSTALL_AGE" != "true" ]]; then
        return 0
    fi

    step "Checking age key"

    if [[ -f "$AGE_KEY_PATH" ]]; then
        log "Age key exists: $AGE_KEY_PATH"
        return 0
    fi

    warn "No age key found at $AGE_KEY_PATH"
    warn "If your dotfiles use encrypted files, create a key:"
    warn "  age-keygen -o $AGE_KEY_PATH"
    warn "Or copy your existing key from another machine"
}

# ── Initialize chezmoi ────────────────────────────────────────────────
init_chezmoi() {
    step "Initializing chezmoi"

    log "Repository: $REPO_URL"

    if [[ -d "$HOME/.local/share/chezmoi" ]]; then
        warn "chezmoi source directory already exists"
        warn "Use 'chezmoi update' to pull latest changes"

        if [[ "$DRY_RUN" == "true" ]]; then
            log "[DRY RUN] Would run: chezmoi update"
        else
            log "Running chezmoi update..."
            chezmoi update -v
        fi
        return 0
    fi

    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY RUN] Would run: chezmoi init --apply $REPO_URL $CHEZMOI_ARGS"
        return 0
    fi

    # shellcheck disable=SC2086
    chezmoi init --apply -v "$REPO_URL" $CHEZMOI_ARGS
    log "chezmoi initialized and applied"
}

# ── Verify ────────────────────────────────────────────────────────────
verify() {
    step "Verification"

    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY RUN] Skipping verification"
        return 0
    fi

    if ! command -v chezmoi > /dev/null 2>&1; then
        err "chezmoi not found in PATH"
        return 1
    fi

    local managed_count
    managed_count="$(chezmoi managed | wc -l)"
    log "Managed files: $managed_count"

    # List managed files
    if [[ "$VERBOSE" == "true" ]]; then
        debug "Managed files:"
        chezmoi managed | while read -r f; do
            debug "  $f"
        done
    fi

    # Check for issues
    local status_count
    status_count="$(chezmoi status | wc -l)"
    if [[ "$status_count" -gt 0 ]]; then
        warn "$status_count files differ from source:"
        chezmoi status
    else
        log "All managed files match source"
    fi

    # Run chezmoi doctor
    debug "Running chezmoi doctor..."
    if [[ "$VERBOSE" == "true" ]]; then
        chezmoi doctor || true
    fi
}

# ── Summary ───────────────────────────────────────────────────────────
summary() {
    step "Summary"

    echo ""
    echo -e "  ${BOLD}Hostname:${RESET}      $HOSTNAME_SHORT"
    echo -e "  ${BOLD}Distro:${RESET}        $DISTRO $DISTRO_VERSION"
    echo -e "  ${BOLD}Repository:${RESET}    $REPO_URL"
    echo -e "  ${BOLD}chezmoi:${RESET}       $(command -v chezmoi 2>/dev/null || echo 'not installed')"

    if command -v age > /dev/null 2>&1; then
        echo -e "  ${BOLD}age:${RESET}           $(age --version)"
    fi

    if [[ -d "$HOME/.local/share/chezmoi" ]]; then
        echo -e "  ${BOLD}Source dir:${RESET}    $HOME/.local/share/chezmoi"
        echo -e "  ${BOLD}Managed:${RESET}       $(chezmoi managed 2>/dev/null | wc -l) files"
    fi

    echo ""

    if [[ "$DRY_RUN" == "true" ]]; then
        echo -e "  ${YELLOW}${BOLD}DRY RUN${RESET} — no changes were made."
        echo -e "  Run with ${BOLD}--force${RESET} to apply."
    elif [[ "$ERRORS" -gt 0 ]]; then
        echo -e "  ${RED}${BOLD}Completed with $ERRORS error(s)${RESET}"
    else
        echo -e "  ${GREEN}${BOLD}Bootstrap complete${RESET}"
    fi

    echo ""
}

# ── Main ──────────────────────────────────────────────────────────────
main() {
    parse_args "$@"
    setup_colors

    echo ""
    echo -e "${BOLD}$SCRIPT_NAME${RESET} — chezmoi dotfile bootstrap"
    echo ""

    if [[ "$DRY_RUN" == "true" ]]; then
        log "Running in ${YELLOW}DRY RUN${RESET} mode (use --force to apply)"
    fi

    detect_system
    install_packages "$PRE_PACKAGES"
    install_chezmoi
    install_age
    setup_age_key
    init_chezmoi
    verify
    summary

    [[ "$ERRORS" -gt 0 ]] && exit 1
    exit 0
}

main "$@"