#!/usr/bin/env bash
#
# tmux-resurrect.sh — snapshot and restore the exact layout of a tmux session
# on Dankclaw, including each pane's start directory and foreground command.
#
# Use it after a reboot (or when the display session is gone) to reopen the
# same windows/panes on the Dankclaw console:
#
#   scripts/tmux-resurrect.sh snapshot      # record the live session to JSON
#   scripts/tmux-resurrect.sh restore       # rebuild it, then attach
#
# Data file: ~/.tmux-resurrect/<session>.json (set TMUX_RES_DIR to relocate).
# Re-run snapshot as the session changes so the restore stays faithful.
#
# How each pane's command is captured: a tmux pane runs one main shell
# (#{pane_pid}). The foreground command is that shell's direct child — chosen
# as the child that launched its own child (ssh, sudo ...) or the first child
# when none did (btop sitting at the top of a pane). An idle shell yields no
# command and is left at its prompt on restore.
set -euo pipefail

TMUX_RES_DIR="${TMUX_RES_DIR:-$HOME/.tmux-resurrect}"
DEFAULT_SESSION="0"

die() { printf 'tmux-resurrect: %s\n' "$*" >&2; exit 1; }
b64() { base64 -w0; }
b64d() { base64 -d; }

# fg_command <shell_pid>: print the foreground command of a tmux pane, or the
# empty string when the shell is idle.
fg_command() {
    local pid="$1" kids=() kid cmd_pid=""
    mapfile -t kids < <(pgrep -P "$pid" 2>/dev/null || true)
    if [ "${#kids[@]}" -eq 0 ]; then return 0; fi
    # Prefer a child that launched its own child (a launcher such as ssh or
    # sudo); fall back to the first direct child (e.g. btop at the top).
    for kid in "${kids[@]}"; do
        if pgrep -P "$kid" >/dev/null 2>&1; then
            cmd_pid="$kid"; break
        fi
    done
    if [ -z "$cmd_pid" ]; then cmd_pid="${kids[0]}"; fi
    tr '\0' ' ' < "/proc/$cmd_pid/cmdline" 2>/dev/null | sed -E 's/[[:space:]]+$//'
}

# snapshot <session> [dest.json]
snapshot() {
    local session="${1:-$DEFAULT_SESSION}" dest="${2:-}"
    local exists
    exists=$(tmux list-sessions -F '#{session_name}' 2>/dev/null | grep -Fxq "$session" && echo 1 || echo 0)
    [ "$exists" = 1 ] || die "no tmux session named '$session'"
    mkdir -p "$TMUX_RES_DIR"
    [ -z "$dest" ] && dest="$TMUX_RES_DIR/$session.json"
    python3 - "$dest" "$session" 3< <(
    {
        while IFS=$'\t' read -r widx wname wlayout; do
            printf 'W\t%s\t%s\t%s\n' "$widx" "$(printf '%s' "$wname" | b64)" "$(printf '%s' "$wlayout" | b64)"
            while IFS=$'\t' read -r pidx ppid pdir; do
                pcmd=$(fg_command "$ppid")
                printf 'P\t%s\t%s\t%s\t%s\n' "$widx" "$pidx" \
                    "$(printf '%s' "$pdir" | b64)" "$(printf '%s' "$pcmd" | b64)"
            done < <(tmux list-panes -t "$session:$widx" \
                        -F $'#{pane_index}\t#{pane_pid}\t#{pane_current_path}' | sort -n -k1,1)
        done < <(tmux list-windows -t "$session" \
                    -F $'#{window_index}\t#{window_name}\t#{window_layout}' | sort -n -k1,1)
    }
    ) > /dev/null <<'PY'
import sys, json, base64, os
dest, session = sys.argv[1], sys.argv[2]
wmap, panes = {}, {}
for line in os.fdopen(3).read().splitlines():
    if not line:
        continue
    p = line.split('\t')
    if p[0] == 'W':
        wmap[int(p[1])] = {"name": base64.b64decode(p[2]).decode(),
                           "layout": base64.b64decode(p[3]).decode()}
    elif p[0] == 'P':
        widx = int(p[1])
        panes.setdefault(widx, []).append({
            "index": int(p[2]),
            "dir": base64.b64decode(p[3]).decode(),
            "cmd": base64.b64decode(p[4]).decode(),
        })
windows = []
for widx in sorted(wmap):
    windows.append({
        "index": widx,
        "name": wmap[widx]["name"],
        "layout": wmap[widx]["layout"],
        "panes": sorted(panes.get(widx, []), key=lambda x: x["index"]),
    })
with open(dest, "w") as f:
    json.dump({"session": session, "windows": windows}, f, indent=2, ensure_ascii=False)
    f.write("\n")
PY
    printf 'snapshot: wrote %s (session %s)\n' "$dest" "$session"
}

# emit_plan <src.json> [session]: python reads the JSON and prints a base64
# TSV plan (W/P records) for the bash orchestrator to consume.
emit_plan() {
    python3 - "$1" "$2" <<'PY'
import sys, json, base64
src, session = sys.argv[1], sys.argv[2]
with open(src) as f:
    data = json.load(f)
wins = data.get("windows", [])
if session:
    wins = [w for w in wins if str(w.get("index")) == "0"]
def b64(s): return base64.b64encode(s.encode()).decode()
for w in wins:
    print('W\t%s\t%s\t%s' % (w["index"], b64(w["name"]), b64(w["layout"])))
    for pan in w["panes"]:
        print('P\t%s\t%s\t%s\t%s' % (w["index"], pan["index"],
                                      b64(pan["dir"]), b64(pan["cmd"])))
PY
}

# restore [session] [--no-attach] [--dry-run]
restore() {
    local session="${1:-$DEFAULT_SESSION}"
    shift || true
    local keep=0 noattach=0 dryrun=0 a
    for a in "$@"; do
        case "$a" in
            --keep) keep=1 ;;
            --no-attach) noattach=1 ;;
            --dry-run) dryrun=1 ;;
            *) die "unknown restore option: $a" ;;
        esac
    done
    local dest="$TMUX_RES_DIR/$session.json"
    [ -f "$dest" ] || die "no snapshot at $dest (run snapshot first)"

    # Restore wrapper: run tmux, or print the command under --dry-run.
    tmuxr() {
        if [ "$dryrun" = 1 ]; then printf '  tmux %s\n' "$*"; else tmux "$@"; fi
    }

    local target_exists
    target_exists=$(tmux list-sessions -F '#{session_name}' 2>/dev/null | grep -Fxq "$session" && echo 1 || echo 0)
    if [ "$target_exists" = 1 ]; then
        if [ "$keep" = 1 ]; then
            die "session '$session' already exists; rebuild without --keep or save a copy"
        fi
        if [ "$dryrun" = 1 ]; then printf '  # kill existing session %s\n' "$session"
        else tmux kill-session --force "$session"; fi
    fi

    # Parse the plan into flat, window-major lists.
    local -a WNAME=() WLAY=() WSTART=() PD=() PC=()
    local tag idx nameb layoutb dirb cmdb
    while IFS=$'\t' read -r tag idx nameb layoutb dirb cmdb; do
        case "$tag" in
            W) WNAME+=("$(printf '%s' "$nameb" | b64d)"); WLAY+=("$(printf '%s' "$layoutb" | b64d)"); WSTART+=("${#PD[@]}") ;;
            P) PD+=("$(printf '%s' "$dirb" | b64d)"); PC+=("$(printf '%s' "$cmdb" | b64d)") ;;
        esac
    done < <(emit_plan "$dest" "")

    local n=${#WNAME[@]} w first last npanes i curwin wname wlayout dir cmd
    for ((w = 0; w < n; w++)); do
        first=${WSTART[w]}
        if (( w + 1 < n )); then last=${WSTART[w + 1]}; else last=${#PD[@]}; fi
        npanes=$(( last - first ))
        wname="${WNAME[w]}"; wlayout="${WLAY[w]}"
        if [ "$dryrun" = 1 ]; then
            if (( w == 0 )); then
                tmuxr new-session -d -s "$session" -n "$wname" -c "${PD[first]}"
            else
                tmuxr new-window -d -t "$session" -n "$wname" -c "${PD[first]}"
            fi
            curwin="$session:$wname"
        elif (( w == 0 )); then
            curwin=$(tmux new-session -d -P -F '#{window_id}' \
                -s "$session" -n "$wname" -c "${PD[first]}")
        else
            curwin=$(tmux new-window -d -P -F '#{window_id}' \
                -t "$session" -n "$wname" -c "${PD[first]}")
        fi
        # Create the remaining panes in recorded order (index == creation order).
        for ((i = first + 1; i < last; i++)); do
            tmuxr split-window -t "$curwin" -c "${PD[i]}"
        done
        # Restore the exact geometry for this window.
        tmuxr select-layout -t "$curwin" "$wlayout"
        # Launch each pane's recorded command, matched to panes by ascending
        # pane index (so this is correct regardless of base-index).
        local -a ordered=()
        mapfile -t ordered < <(tmux list-panes -t "$curwin" \
            -F $'#{pane_id}\t#{pane_target}' | sort -n -k1,1)
        local k=0 tgt
        for tgt in "${ordered[@]}"; do
            cmd="${PC[first + k]}"
            if [ -n "$cmd" ]; then tmuxr send-keys -t "$tgt" "$cmd" Enter; fi
            k=$((k + 1))
        done
    done

    if [ "$noattach" = 1 ]; then
        return 0
    elif [ -t 0 ]; then
        tmuxr attach -t "$session"
    elif [ "$dryrun" = 0 ]; then
        echo "session '$session' rebuilt; attach with: tmux attach -t $session"
    fi
}

usage() {
    cat <<'USAGE'
tmux-resurrect.sh — snapshot & restore a tmux session on Dankclaw.

Usage:
  tmux-resurrect.sh snapshot [session] [dest.json]   record the live session
  tmux-resurrect.sh restore  [session] [options]      rebuild the session
    options:
      --keep        skip if the session already exists (default: rebuild)
      --no-attach   rebuild without attaching to this terminal
      --dry-run     print the commands instead of running them
USAGE
}

main() {
    local cmd="${1:-snapshot}"
    shift || true
    case "$cmd" in
        snapshot) snapshot "$@" ;;
        restore)  restore "$@" ;;
        -h|--help|help) usage ;;
        *) die "unknown command: $cmd (use snapshot|restore)" ;;
    esac
}

main "$@"