#!/usr/bin/bash
# veilor-postinstall — first-login TUI on v0.7+ atomic systems.
#
# Runs ONCE on first SDDM login via the user-mode systemd unit
# `veilor-postinstall.service`. Asks the operator for the small set
# of decisions we deliberately defer from install time:
#   - keyboard / locale
#   - hostname override
#   - GPU drivers (NVIDIA layered via rpm-ostree, mesa = no-op)
#   - package preset (dev / media / homelab — additive, opt-out)
#   - bluetooth opt-in
#   - USBGuard policy snapshot
#   - veilor-doctor first run
# Writes /var/lib/veilor/postinstall-complete on success and disables
# its own autostart unit. Idempotent: safe to re-run.
#
# Style: gum if present, plain bash read fallback. No decorative ASCII.

set -uo pipefail
export TERM="${TERM:-linux}"

STATE_DIR=/var/lib/veilor
DONE_MARKER="$STATE_DIR/postinstall-complete"
LOG=/var/log/veilor-postinstall.log

have() { command -v "$1" >/dev/null 2>&1; }
GUM=$(have gum && echo gum || echo "")

# Always log + tee to stdout for live progress.
mkdir -p "$STATE_DIR" 2>/dev/null || true
exec > >(tee -a "$LOG") 2>&1

if [[ -e $DONE_MARKER && ${1:-} != "--force" ]]; then
    echo "veilor-postinstall already ran (marker: $DONE_MARKER). Pass --force to re-run."
    exit 0
fi

# ── Wrappers ────────────────────────────────────────────────────────
choose() {
    local header=$1; shift
    if [[ -n $GUM ]]; then
        gum choose --header "$header" "$@"
    else
        echo
        echo "$header"
        local i=1
        for opt in "$@"; do printf '  %d) %s\n' "$i" "$opt"; ((i++)); done
        local n
        read -rp "  choice (1-$#): " n
        [[ $n -ge 1 && $n -le $# ]] || return 1
        eval "echo \${$n}"
    fi
}

ask() {
    local prompt=$1 default=${2:-}
    if [[ -n $GUM ]]; then
        gum input --header "$prompt" --value "$default"
    else
        local v
        read -rp "$prompt [$default] " v
        echo "${v:-$default}"
    fi
}

confirm() {
    local prompt=$1
    if [[ -n $GUM ]]; then
        gum confirm "$prompt" && return 0 || return 1
    else
        read -rp "$prompt [y/N] " y
        [[ ${y,,} == y* ]]
    fi
}

say() {
    if [[ -n $GUM ]]; then
        gum style --foreground 212 --bold "$1"
    else
        printf '\n=== %s ===\n' "$1"
    fi
}

# Need root for several actions; re-exec under sudo if not root.
if [[ $EUID -ne 0 ]]; then
    say "veilor-postinstall: sudo required"
    exec sudo -E bash "$0" "$@"
fi

say "veilor-postinstall — one-time setup"
echo "  This runs once. Each step is skippable. Defaults are sane."
echo

# ── 1. Keyboard layout ──────────────────────────────────────────────
KB=$(choose "Keyboard layout" us gb de fr es ru "skip") || KB=skip
if [[ $KB != skip ]]; then
    localectl set-keymap "$KB" 2>/dev/null || true
    echo "  [OK] keymap = $KB"
fi

# ── 2. Locale ───────────────────────────────────────────────────────
LOC=$(choose "Locale" en_US.UTF-8 en_GB.UTF-8 de_DE.UTF-8 fr_FR.UTF-8 "skip") || LOC=skip
if [[ $LOC != skip ]]; then
    localectl set-locale LANG="$LOC" 2>/dev/null || true
    echo "  [OK] locale = $LOC"
fi

# ── 3. Hostname ─────────────────────────────────────────────────────
HN=$(ask "Hostname" "veilor")
if [[ -n $HN && $HN != $(hostnamectl --static 2>/dev/null) ]]; then
    hostnamectl set-hostname "$HN"
    echo "  [OK] hostname = $HN"
fi

# ── 4. GPU drivers ──────────────────────────────────────────────────
GPU=$(choose "GPU drivers" "Skip (use mesa defaults)" "NVIDIA proprietary (akmod-nvidia)" "Intel/AMD mesa (no-op)") || GPU=skip
case "$GPU" in
    *NVIDIA*)
        say "Layering NVIDIA driver — this takes a few minutes"
        rpm-ostree install --idempotent akmod-nvidia xorg-x11-drv-nvidia-cuda \
            && echo "  [OK] NVIDIA driver layered (reboot to use)" \
            || echo "  [WARN] NVIDIA layer failed; check rpm-ostree status"
        ;;
    *) echo "  (skipped GPU layering)" ;;
esac

# ── 5. Package presets (multi-select) ───────────────────────────────
say "Package presets — pick any combination (skip = none)"
PRESET_DEV="git tmux vim-enhanced htop podman skopeo"
PRESET_MEDIA="vlc obs-studio"
PRESET_HOMELAB="wireguard-tools jq yq tmux"

PICKED=()
confirm "Install dev preset? ($PRESET_DEV)"        && PICKED+=($PRESET_DEV) || true
confirm "Install media preset? ($PRESET_MEDIA)"    && PICKED+=($PRESET_MEDIA) || true
confirm "Install homelab preset? ($PRESET_HOMELAB)" && PICKED+=($PRESET_HOMELAB) || true
if (( ${#PICKED[@]} > 0 )); then
    # de-dupe
    UNIQ=$(printf '%s\n' "${PICKED[@]}" | sort -u | tr '\n' ' ')
    say "Layering: $UNIQ"
    rpm-ostree install --idempotent $UNIQ \
        && echo "  [OK] preset packages layered (reboot to use)" \
        || echo "  [WARN] preset layer failed; check rpm-ostree status"
fi

# ── 6. Bluetooth ────────────────────────────────────────────────────
if confirm "Enable Bluetooth?"; then
    systemctl enable --now bluetooth.service 2>/dev/null || true
    echo "  [OK] bluetooth enabled"
else
    echo "  (skipped bluetooth)"
fi

# ── 7. USBGuard snapshot ────────────────────────────────────────────
say "USBGuard policy snapshot"
echo "  Plug in EVERY USB device you trust right now (keyboard,"
echo "  mouse, dock, yubikey, etc.) before continuing."
if confirm "Snapshot current USB devices into the allowlist?"; then
    usbguard generate-policy > /etc/usbguard/rules.conf \
        && echo "  [OK] policy written to /etc/usbguard/rules.conf" \
        || echo "  [WARN] generate-policy failed"
    systemctl restart usbguard 2>/dev/null || true
fi

# ── 8. veilor-doctor ────────────────────────────────────────────────
if confirm "Run veilor-doctor now?"; then
    veilor-doctor || true
fi

# ── Done ────────────────────────────────────────────────────────────
date -u +"%Y-%m-%dT%H:%M:%SZ" > "$DONE_MARKER"
say "veilor-postinstall complete"
echo "  Marker written: $DONE_MARKER"
echo "  Disabling autostart unit so this never runs again."
systemctl --user --global disable veilor-postinstall.service 2>/dev/null || true
systemctl disable veilor-postinstall.service 2>/dev/null || true
echo
echo "  If you layered any packages or drivers, reboot to activate."