#!/usr/bin/env bash
# SPDX-License-Identifier: GPL-3.0-only
# eniatos-machine - systemd-nspawn containers with read-only ZFS rootfs built from a recipe,
# updated by the host: snapshot -> clone -> pacman -Syu in the clone -> swap -> health check
set -uo pipefail

PROG=eniatos-machine
CONF=/etc/eniatos/machine.conf
RECIPES=/etc/eniatos/machines
STATE=/var/lib/eniatos-machine
MOTD=/etc/motd.d/51-eniatos-machine
UPDATE_STATE=/var/lib/eniatos-update

POOL=$(tr -d '[:space:]' < /etc/hostname 2>/dev/null); POOL=${POOL,,}   # EniatOS: pool = hostname
MACHINES_DS=$POOL/var/lib_machines       # rootfs datasets: $MACHINES_DS/<name>
STATE_DS=$POOL/srv/machines              # persistent data: $STATE_DS/<name>
STATE_MNT=/srv/machines                  # mountpoint of $STATE_DS
KEEP=10                                  # previous rootfs versions kept per machine
HEALTH_IP=""                             # public address used by health checks
HEALTH_TIMEOUT=90
# shellcheck source=/dev/null
[[ -r $CONF ]] && . "$CONF"

usage() {
    cat <<'EOF'
Usage: eniatos-machine build NAME        build rootfs from the recipe (first install or full rebuild)
       eniatos-machine update [NAME...]  update rootfs of the given / all machines (timer)
       eniatos-machine rollback NAME [VERSION]  switch back to a previous rootfs (default: newest)
       eniatos-machine list [NAME]       machines and kept rootfs versions
       eniatos-machine health NAME       run the health check
       eniatos-machine prune NAME        drop rootfs versions beyond KEEP
       eniatos-machine status            last update report
Recipe:  /etc/eniatos/machines/NAME/  (machine.conf packages enable NAME.nspawn overlay-pre/ overlay/ seed/ service.conf)
Config:  /etc/eniatos/machine.conf
EOF
    exit "${1:-0}"
}

log() { echo "$*"; logger -t "$PROG" -p user.info -- "$*" 2>/dev/null || true; }
warn() { echo "WARNING: $*" >&2; logger -t "$PROG" -p user.warning -- "$*" 2>/dev/null || true; }
err() { echo "ERROR: $*" >&2; logger -t "$PROG" -p user.err -- "$*" 2>/dev/null || true; }
die() { err "$*"; exit 1; }

mnt() { zfs get -H -o value mountpoint "$1"; }
exists() { zfs list -H -o name "$1" >/dev/null 2>&1; }
unit() { echo "systemd-nspawn@$1.service"; }

# load_recipe NAME - recipe variables: PACKAGES ENABLE STATE_DIRS HEALTH_UNITS HEALTH_HOSTS
load_recipe() {
    NAME=$1 R=$RECIPES/$1
    [[ $NAME =~ ^[a-z0-9][a-z0-9-]*$ ]] || die "invalid machine name: $NAME"
    [[ -d $R ]] || die "no recipe: $R"
    STATE_DIRS="" HEALTH_UNITS="" HEALTH_HOSTS=""
    # shellcheck source=/dev/null
    [[ -r $R/machine.conf ]] && . "$R/machine.conf"
    DS=$MACHINES_DS/$NAME
}

all_machines() {
    local d
    for d in "$RECIPES"/*/; do d=${d%/}; [[ -d $d ]] && echo "${d##*/}"; done
}

# --- status / motd ---------------------------------------------------------
declare -A REPORT=()
report() { REPORT[$1]=$2; }
write_status() {
    local m bad=() status=ok
    mkdir -p "$STATE"
    {
        echo "time=$(date '+%Y-%m-%d %H:%M %Z')"
        for m in "${!REPORT[@]}"; do
            echo "$m: ${REPORT[$m]}"
            [[ ${REPORT[$m]} == ok* || ${REPORT[$m]} == unchanged* ]] || bad+=("$m: ${REPORT[$m]}")
        done
    } > "$STATE/status"
    local pacnew
    pacnew=$(find /var/lib/machines/*/etc -maxdepth 3 -name '*.pacnew' 2>/dev/null | grep -v '/[^/]*\.[0-9-]*/etc/' | sort)
    if [[ ${#bad[@]} -gt 0 || -n $pacnew ]]; then
        mkdir -p "${MOTD%/*}"
        {
            echo
            [[ ${#bad[@]} -gt 0 ]] && printf 'eniatos-machine: %s\n' "${bad[@]}"
            [[ -n $pacnew ]] && { echo "eniatos-machine: .pacnew in container images (update the recipe overlay):"; echo "$pacnew"; }
            echo "Details: eniatos-machine status, journalctl -t $PROG"
            echo
        } > "$MOTD"
        [[ ${#bad[@]} -gt 0 ]] && status=failed
    else
        rm -f "$MOTD"
    fi
    [[ $status == ok ]]
}

# --- rootfs construction ---------------------------------------------------
copy_tree() {  # copy_tree SRC DST - files owned by root, modes kept
    [[ -d $1 ]] || return 0
    cp -a --no-preserve=ownership "$1/." "$2/"
}

# finalize ROOT - deterministic image customisation after packages are installed
finalize() {
    local root=$1 id
    # pacman setup of the host (CachyOS repos), so updates inside keep the same repositories
    install -m644 /etc/pacman.conf "$root/etc/pacman.conf"
    cp -a /etc/pacman.d/*mirrorlist "$root/etc/pacman.d/"
    copy_tree "$R/overlay" "$root"
    if [[ ! -s $R/machine-id ]]; then
        systemd-id128 new > "$R/machine-id"
        log "$NAME: generated machine-id $(< "$R/machine-id")"
    fi
    install -m444 "$R/machine-id" "$root/etc/machine-id"
    if [[ -s $R/enable ]]; then
        # shellcheck disable=SC2046
        systemctl --root="$root" enable $(grep -vE '^\s*(#|$)' "$R/enable") 2>&1 | grep -v '^Created symlink' || true
    fi
    if grep -qx 'systemd-resolved.service\|systemd-resolved' "$R/enable" 2>/dev/null; then
        ln -sfn ../run/systemd/resolve/stub-resolv.conf "$root/etc/resolv.conf"
    fi
    # bind mount targets must exist - the rootfs is read-only at runtime
    local target
    for target in $(sed -nE 's/^(Bind|BindReadOnly)=[^:]+:([^:]+).*/\2/p' "$R/$NAME.nspawn"); do
        [[ $target == /var/* ]] || mkdir -p "$root$target"
    done
    rm -rf "$root"/var/cache/pacman/pkg/* "$root"/var/log/pacman.log
    id=$(date +%s); echo "$id" > "$root/etc/eniatos-machine-build"
}

install_units() {
    install -Dm644 "$R/$NAME.nspawn" "/etc/systemd/nspawn/$NAME.nspawn"
    if [[ -s $R/service.conf ]]; then
        install -Dm644 "$R/service.conf" "/etc/systemd/system/$(unit "$NAME").d/eniatos.conf"
    fi
    systemctl daemon-reload
}

ensure_state() {
    local spec dir uid gid mode base
    exists "$STATE_DS" || zfs create -p -o mountpoint="$STATE_MNT" "$STATE_DS"
    exists "$STATE_DS/$NAME" || zfs create "$STATE_DS/$NAME"
    base=$(mnt "$STATE_DS/$NAME")
    for spec in $STATE_DIRS; do   # dir:uid:gid:mode
        IFS=: read -r dir uid gid mode <<< "$spec"
        if [[ ! -d $base/$dir ]]; then
            install -d -o "$uid" -g "$gid" -m "$mode" "$base/$dir"
            [[ -d $R/seed/$dir ]] && { cp -a --no-preserve=ownership "$R/seed/$dir/." "$base/$dir/"; chown -R "$uid:$gid" "$base/$dir"; }
            log "$NAME: created state $base/$dir"
        fi
    done
}

do_build() {
    load_recipe "$1"
    [[ -s $R/packages && -s $R/$NAME.nspawn ]] || die "$NAME: recipe needs packages and $NAME.nspawn"
    local new=$DS.new root ts
    ts=$(date +%Y%m%d-%H%M%S)
    exists "$new" && { zfs destroy -r "$new" || die "cannot remove stale $new"; }
    zfs create "$new" || die "zfs create $new failed"
    root=$(mnt "$new")
    copy_tree "$R/overlay-pre" "$root"
    log "$NAME: pacstrap into $root"
    # shellcheck disable=SC2046
    if ! pacstrap -c "$root" $(grep -vE '^\s*(#|$)' "$R/packages"); then
        zfs destroy -r "$new"; die "$NAME: pacstrap failed"
    fi
    finalize "$root"
    zfs set readonly=on "$new"
    ensure_state
    install_units
    if exists "$DS"; then
        swap_in "$ts" || die "$NAME: rebuilt rootfs failed the health check, rolled back"
    else
        zfs rename "$new" "$DS" || die "rename $new failed"
        systemctl enable "$(unit "$NAME")" machines.target >/dev/null 2>&1
        systemctl start "$(unit "$NAME")" || die "$NAME: start failed"
        health "$NAME" || warn "$NAME: health check failed after first build"
    fi
    log "$NAME: built"
}

# --- update ---------------------------------------------------------------
pkgs_of() { ls "$1/var/lib/pacman/local" 2>/dev/null; }

# update_one NAME - returns 0 ok/unchanged, 1 failure
update_one() {
    load_recipe "$1"
    exists "$DS" || { report "$NAME" "failed: no rootfs, run build"; return 1; }
    local ts snap new=$DS.new root before after out rc
    ts=$(date +%Y%m%d-%H%M%S)
    snap=$DS@upd-$ts
    exists "$new" && zfs destroy -r "$new"
    zfs snapshot "$snap" || { report "$NAME" "failed: snapshot"; return 1; }
    zfs clone -o readonly=off "$snap" "$new" || { zfs destroy "$snap"; report "$NAME" "failed: clone"; return 1; }
    root=$(mnt "$new")
    before=$(pkgs_of "$root")
    log "$NAME: pacman -Syu in $root"
    out=$(systemd-nspawn -q -D "$root" --machine="$NAME-upd" --register=no --private-users=no \
            --bind=/var/cache/pacman/pkg --bind-ro=/run/systemd/resolve --resolv-conf=off \
            --setenv=LANG=C.UTF-8 \
            sh -c 'pacman -Sy --needed --noconfirm archlinux-keyring cachyos-keyring && pacman -Su --noconfirm' 2>&1)
    rc=$?
    echo "$out"
    if [[ $rc -ne 0 ]]; then
        zfs destroy -r "$new"; zfs destroy "$snap"
        report "$NAME" "failed: pacman (rc=$rc), running rootfs unchanged"
        return 1
    fi
    after=$(pkgs_of "$root")
    if [[ $before == "$after" ]]; then
        zfs destroy -r "$new"; zfs destroy "$snap"
        log "$NAME: nothing to update"
        report "$NAME" "unchanged"
        return 0
    fi
    rm -rf "$root"/var/cache/pacman/pkg/*
    zfs set readonly=on "$new"
    if swap_in "$ts"; then
        report "$NAME" "ok: $(comm -13 <(echo "$before") <(echo "$after") | wc -l) packages updated"
        prune "$NAME"
        return 0
    fi
    report "$NAME" "failed: health check after update, rolled back to previous rootfs"
    return 1
}

# swap_in TS - $DS.new becomes $DS, the old one $DS.TS; health check, automatic rollback
swap_in() {
    local ts=$1 old=$DS.$1 u
    u=$(unit "$NAME")
    log "$NAME: switching rootfs"
    systemctl stop "$u"
    zfs rename "$DS" "$old" || { systemctl start "$u"; return 1; }
    if ! zfs rename "$DS.new" "$DS"; then
        zfs rename "$old" "$DS"; systemctl start "$u"; return 1
    fi
    [[ $(zfs get -H -o value origin "$DS") != - ]] && zfs promote "$DS"
    systemctl start "$u"
    if health "$NAME"; then
        log "$NAME: running new rootfs (previous kept as ${old##*/})"
        return 0
    fi
    err "$NAME: health check failed, rolling back to ${old##*/}"
    systemctl stop "$u"
    zfs rename "$DS" "$DS.failed-$ts"
    zfs rename "$old" "$DS"
    [[ $(zfs get -H -o value origin "$DS") != - ]] && zfs promote "$DS"
    systemctl start "$u"
    health "$NAME" || err "$NAME: still unhealthy after rollback"
    return 1
}

# --- health ---------------------------------------------------------------
health() {
    load_recipe "$1"
    local deadline state u h code ok
    deadline=$(( $(date +%s) + HEALTH_TIMEOUT ))
    while :; do
        state=$(timeout 10 systemctl -M "$NAME" is-system-running 2>/dev/null)
        [[ $state == running ]] && break
        if (( $(date +%s) >= deadline )); then
            err "$NAME: container state '${state:-unreachable}'"
            [[ $state == degraded ]] && timeout 10 systemctl -M "$NAME" --failed --no-legend --plain >&2
            return 1
        fi
        sleep 2
    done
    for u in $HEALTH_UNITS; do
        systemctl -M "$NAME" is-active -q "$u" || { err "$NAME: $u not active"; return 1; }
    done
    for h in $HEALTH_HOSTS; do
        ok=0
        while (( $(date +%s) < deadline )); do
            # plain HTTP through the proxy: exercises proxy -> PROXY protocol -> Caddy,
            # independent of certificate state (expects the HTTPS redirect or a page)
            code=$(curl -sS -o /dev/null -m 10 -w '%{http_code}' ${HEALTH_IP:+--resolve "$h:80:$HEALTH_IP"} "http://$h/" 2>/dev/null)
            [[ $code =~ ^[23] ]] && { ok=1; break; }
            sleep 2
        done
        [[ $ok == 1 ]] || { err "$NAME: http://$h/ -> ${code:-no answer}"; return 1; }
    done
    log "$NAME: healthy"
}

# --- retention / rollback -------------------------------------------------
versions() {  # previous rootfs datasets, newest first
    zfs list -H -o name -s creation -d 1 "$MACHINES_DS" 2>/dev/null \
        | grep -E "^$DS\.(failed-)?[0-9]{8}-[0-9]{6}$" | tac
}

prune() {
    load_recipe "$1"
    local v s n=0
    for v in $(versions); do
        n=$((n + 1))
        (( n > KEEP )) || continue
        zfs destroy "$v" && log "$NAME: pruned ${v##*/}" || warn "$NAME: cannot destroy $v"
    done
    # snapshots no longer origins of any kept version
    for s in $(zfs list -H -t snapshot -o name "$DS" 2>/dev/null | grep "@upd-"); do
        [[ $(zfs get -H -o value clones "$s") == - || -z $(zfs get -H -o value clones "$s") ]] && zfs destroy "$s"
    done
}

do_rollback() {
    load_recipe "$1"
    local target=${2:-} ts u
    if [[ -z $target ]]; then
        target=$(versions | grep -v failed | head -1)
    else
        [[ $target == */* ]] || target=$DS.${target#"$NAME".}
    fi
    [[ -n $target ]] && exists "$target" || die "$NAME: no such version ${target:-}"
    ts=$(date +%Y%m%d-%H%M%S); u=$(unit "$NAME")
    log "$NAME: rollback to ${target##*/}"
    systemctl stop "$u"
    zfs rename "$DS" "$DS.$ts" || { systemctl start "$u"; die "rename failed"; }
    zfs rename "$target" "$DS" || { zfs rename "$DS.$ts" "$DS"; systemctl start "$u"; die "rename failed"; }
    [[ $(zfs get -H -o value origin "$DS") != - ]] && zfs promote "$DS"
    systemctl start "$u"
    health "$NAME"
}

do_list() {
    local m v
    for m in ${1:-$(all_machines)}; do
        load_recipe "$m"
        printf '%-16s %-10s %s\n' "$NAME" "$(systemctl is-active "$(unit "$NAME")")" \
            "$(exists "$DS" && zfs get -H -o value used,creation "$DS" | xargs || echo 'not built')"
        for v in $(versions); do
            printf '  %-36s %s\n' "${v##*/}" "$(zfs get -H -o value creation "$v")"
        done
    done
}

do_update() {
    exec 9> /run/eniatos-machine.lock
    flock -n 9 || die "already running"
    log "=== eniatos-machine update $(date '+%F %T %Z') ==="
    local m machines=("$@") rc=0 health_out
    [[ ${#machines[@]} -gt 0 ]] || mapfile -t machines < <(all_machines)
    health_out=$(zpool status -x 2>&1)
    if [[ $health_out != "all pools are healthy" ]]; then
        for m in "${machines[@]}"; do report "$m" "skipped: zpool not healthy"; done
        write_status; exit 1
    fi
    if [[ -e $UPDATE_STATE/news-pending ]]; then
        for m in "${machines[@]}"; do report "$m" "held: unread Arch News (eniatos-update --status)"; done
        write_status; exit 0
    fi
    for m in "${machines[@]}"; do
        update_one "$m" || rc=1
    done
    write_status || rc=1
    exit $rc
}

[[ $(id -u) -eq 0 ]] || { echo "must be root" >&2; exit 1; }
case "${1:-}" in
    build) [[ -n ${2:-} ]] || usage 1; do_build "$2" ;;
    update) shift; do_update "$@" ;;
    rollback) [[ -n ${2:-} ]] || usage 1; do_rollback "$2" "${3:-}" ;;
    list) do_list "${2:-}" ;;
    health) [[ -n ${2:-} ]] || usage 1; health "$2" ;;
    prune) [[ -n ${2:-} ]] || usage 1; prune "$2" ;;
    status) cat "$STATE/status" 2>/dev/null || echo "no runs yet"; cat "$MOTD" 2>/dev/null ;;
    -h|--help) usage 0 ;;
    *) usage 1 ;;
esac
