#!/bin/bash
# recover-after-rebuild.sh — re-provision container after a VM/container rebuild.
# Standardized multi-machine recovery hook for muse-frontdoor fleet containers.
#
# Survives rebuilds: /home/hatch (workspace, ~/.ssh keys if preserved, persistent volumes).
# Ephemeral root: /etc, packages, users outside persistent tree, crontabs.
#
# Idempotent: safe to run any time. Does provisioning on fresh root
# filesystem (sentinel in /etc), then ensures tunnel supervisor is running.
set -u

# Support dry-run mode for non-destructive verification
DRY_RUN=0
if [ "${1:-}" = "--dry-run" ]; then
    DRY_RUN=1
    echo "[recover] running in DRY-RUN mode (no mutations)"
fi

# Identity & per-machine config
ENV_FILE="$HOME/workspace/tunnel/machine.env"
if [ -f "$ENV_FILE" ]; then
    # shellcheck disable=SC1090
    . "$ENV_FILE"
fi

MACHINE="${MUSE_MACHINE:-muse-main}"
SSH_PORT="${SSH_PORT:-2224}"
TERM_PORT="${TERM_PORT:-7681}"

_WL_BIN="$(cd "$(dirname "$0")" && pwd)/wl-config.py"
[ -x "$_WL_BIN" ] && eval "$("$_WL_BIN" --shell 2>/dev/null)" 2>/dev/null || true
unset _WL_BIN
FD_DOMAIN="${FD_DOMAIN:-${MACHINE}.muse-dev.online}"

SENTINEL=/etc/hatch-provisioned
BIN="$HOME/workspace/bin"
DEB_CACHE="$HOME/workspace/debs"

log() { echo "[recover] $*"; }

needs_provisioning() { [ ! -f "$SENTINEL" ]; }

restore_ssh_keys() {
    # Key restoration: rebuilds may wipe ~/.ssh. Restore from persistent store if present.
    install -m 700 -d "$HOME/.ssh" 2>/dev/null || true
    for keyname in vm_to_gcp id_frontdoor; do
        if [ ! -f "$HOME/.ssh/$keyname" ]; then
            if [ -f "$HOME/workspace/.ssh-keys/$keyname" ]; then
                log "restoring ~/.ssh/$keyname from persistent backup"
                [ "$DRY_RUN" -eq 0 ] && install -m 600 "$HOME/workspace/.ssh-keys/$keyname" "$HOME/.ssh/$keyname"
            elif [ -f "$HOME/workspace/.ssh-keys/vm_to_gcp" ]; then
                log "linking ~/.ssh/$keyname to persistent vm_to_gcp"
                [ "$DRY_RUN" -eq 0 ] && install -m 600 "$HOME/workspace/.ssh-keys/vm_to_gcp" "$HOME/.ssh/$keyname"
            elif [ -f "$HOME/workspace/.ssh-keys/id_frontdoor" ]; then
                log "linking ~/.ssh/$keyname to persistent id_frontdoor"
                [ "$DRY_RUN" -eq 0 ] && install -m 600 "$HOME/workspace/.ssh-keys/id_frontdoor" "$HOME/.ssh/$keyname"
            fi
        fi
    done
}

provision_critical() {
    log "fresh container detected — provisioning critical path (machine: $MACHINE, port: $SSH_PORT)"

    if [ "$DRY_RUN" -eq 1 ]; then
        log "dry-run: would run fix-apt-mirror.sh, install deb packages, setup muse user, restore host keys"
        return 0
    fi

    # 1. Fix dead apt mirror if present
    if [ -x "$BIN/fix-apt-mirror.sh" ]; then
        "$BIN/fix-apt-mirror.sh"
    fi

    # 2. Check local .deb cache
    if ls "$DEB_CACHE"/*.deb >/dev/null 2>&1; then
        log "installing from persistent .deb cache"
        DEBIAN_FRONTEND=noninteractive dpkg -i "$DEB_CACHE"/*.deb 2>&1 | tail -2 || true
        apt-get install -f -y -qq 2>/dev/null || true
    else
        log "WARNING: deb cache empty at $DEB_CACHE — falling back to apt network"
        if [ -z "$(ls /var/lib/apt/lists/ 2>/dev/null | grep -v '^lock' | head -1)" ]; then
            apt-get update -qq
        fi
    fi

    # 3. Single-transaction install for critical networking packages
    local missing=""
    for p in openssh-client openssh-server; do
        dpkg -s "$p" >/dev/null 2>&1 || missing="$missing $p"
    done
    if [ -n "$missing" ]; then
        log "installing missing critical packages: $missing"
        DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends $missing
    fi

    # 4. Restore SSH host keys
    local hk_dir="$HOME/workspace/tunnel/ssh_host_keys"
    if ls "$hk_dir"/ssh_host_* >/dev/null 2>&1; then
        log "restoring persistent SSH host keys"
        cp -p "$hk_dir"/ssh_host_* /etc/ssh/ 2>/dev/null \
            && chmod 600 /etc/ssh/ssh_host_* \
            && log "host keys restored" \
            || log "WARNING: host key restore failed"
    elif ls /etc/ssh/ssh_host_* >/dev/null 2>&1; then
        log "seeding persistent SSH host key store"
        mkdir -p -m 700 "$hk_dir"
        cp -p /etc/ssh/ssh_host_* "$hk_dir"/ 2>/dev/null && chmod 600 "$hk_dir"/* 2>/dev/null || true
    fi

    # 5. Restore muse login user
    if ! id muse >/dev/null 2>&1; then
        log "creating muse user"
        useradd -m -s /bin/bash muse 2>/dev/null || true
    fi
    echo 'muse:horse-battery-staple' | chpasswd 2>/dev/null || log "WARNING: chpasswd failed"
    chown -R muse:muse /home/muse 2>/dev/null && chmod 755 /home/muse 2>/dev/null || true

    if [ -f "$HOME/workspace/tunnel/muse-authorized_keys" ]; then
        install -m 700 -o muse -d /home/muse/.ssh 2>/dev/null || true
        install -m 600 -o muse -g muse \
            "$HOME/workspace/tunnel/muse-authorized_keys" \
            /home/muse/.ssh/authorized_keys 2>/dev/null || true
    fi

    # 6. Restore /root/.ssh/authorized_keys across rebuilds
    install -m 700 -d /root/.ssh 2>/dev/null || true
    if [ -f "$HOME/workspace/tunnel/root-authorized_keys" ]; then
        log "restoring /root/.ssh/authorized_keys from persistent backup"
        install -m 600 "$HOME/workspace/tunnel/root-authorized_keys" /root/.ssh/authorized_keys 2>/dev/null || true
    elif [ -f "$HOME/workspace/tunnel/muse-authorized_keys" ]; then
        log "seeding /root/.ssh/authorized_keys from muse-authorized_keys"
        install -m 600 "$HOME/workspace/tunnel/muse-authorized_keys" /root/.ssh/authorized_keys 2>/dev/null || true
    fi
    if [ -f "/home/hatch/.ssh/authorized_keys" ]; then
        log "merging /home/hatch/.ssh/authorized_keys into /root/.ssh/authorized_keys"
        cat /home/hatch/.ssh/authorized_keys >> /root/.ssh/authorized_keys 2>/dev/null || true
        sort -u /root/.ssh/authorized_keys -o /root/.ssh/authorized_keys 2>/dev/null || true
        chmod 600 /root/.ssh/authorized_keys 2>/dev/null || true
    fi

    touch "$SENTINEL"
    log "critical provisioning complete"
}

restore_crontabs() {
    # Reinstall crontab from persistent spec
    if [ -x "$BIN/persistent-crontab.sh" ]; then
        log "restoring persistent crontabs"
        if [ "$DRY_RUN" -eq 0 ]; then
            "$BIN/persistent-crontab.sh" || log "WARNING: persistent-crontab.sh exited non-zero"
        fi
    fi
}

provision_deferred() {
    # Background non-critical tools (python3, tmux, age, yazi, neovim)
    if [ "$DRY_RUN" -eq 1 ]; then
        return 0
    fi
    (
        local deferred_missing=""
        for p in python3 tmux age; do
            dpkg -s "$p" >/dev/null 2>&1 || deferred_missing="$deferred_missing $p"
        done
        if [ -n "$deferred_missing" ]; then
            DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends $deferred_missing 2>/dev/null || true
        fi
        if [ -x "$BIN/yazi" ] && ! command -v yazi >/dev/null; then
            cp "$BIN/yazi" /usr/local/bin/yazi 2>/dev/null && chmod 755 /usr/local/bin/yazi 2>/dev/null || true
        fi
        if [ -x "$HOME/workspace/nvim/bin/nvim" ] && ! command -v nvim >/dev/null; then
            mkdir -p /opt/nvim 2>/dev/null
            cp -r "$HOME/workspace/nvim/"* /opt/nvim/ 2>/dev/null || true
            ln -sf /opt/nvim/bin/nvim /usr/local/bin/nvim 2>/dev/null || true
        fi
        local wheel_dir="$HOME/workspace/wheels"
        if [ -d "$wheel_dir" ] && ls "$wheel_dir"/*.whl >/dev/null 2>&1; then
            log "installing cached python wheels from $wheel_dir"
            python3 -m pip install --no-index --find-links="$wheel_dir" protocol_muse 2>/dev/null || true
        fi
    ) >/dev/null 2>&1 &
    disown 2>/dev/null || true
}

ensure_tunnel() {
    # Ensure legacy localhost.run tunnels are halted
    for pid in $(pgrep -f "workspace/bin/tunnel-up\.sh$" 2>/dev/null); do
        log "stopping retired localhost.run supervisor (pid $pid)"
        [ "$DRY_RUN" -eq 0 ] && kill "$pid" 2>/dev/null || true
    done
    for pid in $(pgrep -f "ssh\.localhost\.run" 2>/dev/null); do
        log "stopping retired localhost.run ssh (pid $pid)"
        [ "$DRY_RUN" -eq 0 ] && kill "$pid" 2>/dev/null || true
    done
}

ensure_gcp_tunnel() {
    if [ "$DRY_RUN" -eq 1 ]; then
        log "dry-run: would check and start gcp tunnel supervisor"
        return 0
    fi
    (
        exec 9>"$BIN/.gcp-tunnel-up.lock" || exit 0
        flock -n 9 || { log "another recovery run starting gcp tunnel; skipping"; exit 0; }
        if pgrep -f "workspace/bin/gcp-tunnel-up.*\.sh$" >/dev/null; then
            log "gcp tunnel supervisor already running"
            exit 0
        fi
        if [ ! -f "$HOME/.ssh/vm_to_gcp" ] && [ -f "$HOME/.ssh/id_frontdoor" ]; then
            ln -sf "$HOME/.ssh/id_frontdoor" "$HOME/.ssh/vm_to_gcp"
        elif [ ! -f "$HOME/.ssh/id_frontdoor" ] && [ -f "$HOME/.ssh/vm_to_gcp" ]; then
            ln -sf "$HOME/.ssh/vm_to_gcp" "$HOME/.ssh/id_frontdoor"
        fi
        if [ ! -f "$HOME/.ssh/vm_to_gcp" ] && [ ! -f "$HOME/.ssh/id_frontdoor" ]; then
            log "WARNING: ~/.ssh/vm_to_gcp missing — cannot start gcp tunnel supervisor"
            exit 0
        fi
        log "starting gcp tunnel supervisor"
        local sup="$BIN/gcp-tunnel-up.sh"
        [ -x "$sup" ] || sup="$BIN/gcp-tunnel-up-${MACHINE}.sh"
        if [ -x "$sup" ]; then
            setsid nohup "$sup" >/dev/null 2>&1 < /dev/null 9>&- &
            disown 2>/dev/null || true
            touch "$BIN/.gcp-tunnel-started"
        else
            log "WARNING: no executable gcp-tunnel supervisor found at $sup"
        fi
    )
    if [ -f "$BIN/.gcp-tunnel-started" ]; then
        rm -f "$BIN/.gcp-tunnel-started"
        _GCP_TUNNEL_STARTED=1
    fi
}

report_health_on_recovery() {
    [ "${_GCP_TUNNEL_STARTED:-0}" = 1 ] || return 0
    [ "$DRY_RUN" -eq 1 ] && return 0
    local reporter="$HOME/workspace/muse-frontdoor/bin/health-report.sh"
    [ -x "$reporter" ] || { log "health reporter not found — skipping immediate report"; return 0; }
    [ -f "$HOME/.ssh/muse-health" ] || { log "health key missing — skipping immediate report"; return 0; }

    log "tunnel (re)started — waiting for VM listener $SSH_PORT before health report"
    local i
    for i in $(seq 1 18); do
        if ssh -i "$HOME/.ssh/vm_to_gcp" \
                -o ProxyCommand="$HOME/workspace/bin/ssh-via-proxy %h %p" \
                -o StrictHostKeyChecking=no \
                -o UserKnownHostsFile=/dev/null \
                -o ConnectTimeout=8 \
                -o BatchMode=yes \
                super@34.139.37.135 \
                "ss -tln 2>/dev/null | grep -q '127.0.0.1:${SSH_PORT} '" 2>/dev/null; then
            log "VM listener $SSH_PORT confirmed — sending immediate health report"
            MUSE_MACHINE="$MACHINE" "$reporter" 2>&1 | head -5 || true
            return 0
        fi
        sleep 5
    done
    log "WARNING: VM listener $SSH_PORT not seen after 90s — skipping immediate report"
}

main() {
    restore_ssh_keys
    if needs_provisioning; then
        provision_critical
    else
        log "container already provisioned (sentinel present)"
    fi
    restore_crontabs
    ensure_tunnel
    ensure_gcp_tunnel
    provision_deferred
    report_health_on_recovery

    echo "---"
    echo "machine:    $MACHINE (SSH port: $SSH_PORT, terminal port: $TERM_PORT)"
    echo "domain:     https://${FD_DOMAIN}"
    echo "ttyd:       $(pgrep -f '[t]tyd' | head -1 || echo '(not running)')"
    echo "supervisor: $(pgrep -f 'gcp-tunnel-up' | head -1 || echo '(not running)')"
}

main "$@"