#!/usr/bin/env bash

# Tracks the state of Claude Code and pi sessions, keyed by the tmux pane each
# one runs in, so ,agents can rank them by who is waiting on a human.
#
# States, kind, and app follow the vocabulary of OSC 7501, the Program Status
# Protocol (https://www.superlogical.com/rex/docs/build/program-status), so a
# producer here can emit the sequence later without translating. Nothing in
# this stack consumes it yet, which is why the record still lives here.
#
# Written by ai/skills/agent-state/hooks/register.ts and by the pi extension
# ai/pi/extensions/agent-state.ts; read by ai/bin/,agents and by the sketchybar
# agents plugin. Those two readers also write, through `visited`: jumping to a
# blocked session records the visit so the picker can move on to the next one
# instead of offering the same session again, and jumping to a done session
# marks it idle.
#
# `list` is not read-only: it prunes rows whose pane is gone. The sketchybar
# plugin runs it every time the hook reports a state change, so most pruning
# happens there rather than when a human opens the picker. Nothing polls, so a
# pane that dies without a SessionEnd stays listed until the next hook fires.
#
# State lives in a SQLite database, one row per pane. Hooks fire concurrently
# across sessions, so the database opens in WAL mode with a busy timeout:
# writers don't block the reader, and a writer that collides with another
# retries instead of failing. Every value is interpolated into SQL through
# sql_quote, never pasted raw.

set -euo pipefail
if [[ "${TRACE-0}" == "1" ]]; then
    set -o xtrace
fi

STATE_DB="${AGENT_STATE_DB:-${XDG_STATE_HOME:-$HOME/.local/state}/agent-state/state.db}"

# Applied to every connection. WAL and busy_timeout are the two that matter for
# concurrent hooks; synchronous=NORMAL is safe under WAL and keeps a hook from
# waiting on an fsync inside Claude's tool loop. journal_mode is persistent in
# the file, but busy_timeout is per-connection, so this runs every time.
#
# The .output fencing is load-bearing: an assigning PRAGMA *returns a row*, so
# without it `journal_mode = WAL` prints "wal" onto stdout, which corrupts the
# tab-separated rows `list` emits and makes the hook write to stdout inside
# Claude's tool loop.
# busy_timeout MUST come first. Switching journal_mode takes an exclusive lock,
# and with the default zero timeout a writer that collides with another fails
# instantly with "database is locked" rather than waiting — which showed up as
# roughly 1 in 15 dropped writes under concurrent hooks.
readonly PRAGMAS="
.output /dev/null
PRAGMA busy_timeout = 3000;
PRAGMA journal_mode = WAL;
PRAGMA synchronous = NORMAL;
.output stdout
"

# Rows are rebuilt by the next hook to fire, so a schema change renames the
# table and drops the old one rather than migrating it.
readonly SCHEMA="
DROP TABLE IF EXISTS panes;
CREATE TABLE IF NOT EXISTS records (
  pane       TEXT PRIMARY KEY,
  state      TEXT NOT NULL,
  kind       TEXT NOT NULL DEFAULT '',
  app        TEXT NOT NULL DEFAULT '',
  updated_at INTEGER NOT NULL,
  cwd        TEXT NOT NULL DEFAULT '',
  message    TEXT NOT NULL DEFAULT ''
);
-- Its own table rather than a column on records, so a records schema change
-- drops nothing a hook cannot rebuild.
CREATE TABLE IF NOT EXISTS visits (
  pane       TEXT PRIMARY KEY,
  visited_at INTEGER NOT NULL
);
"

# Renders a value as a SQL string literal. SQLite escapes a single quote by
# doubling it, and that is the only escape a single-quoted literal needs.
sql_quote() {
    printf "'%s'" "${1//\'/\'\'}"
}

db() {
    mkdir -p "$(dirname "$STATE_DB")"
    # -batch stops sqlite3 from ever prompting; -noheader keeps output parseable.
    sqlite3 -batch -noheader "$STATE_DB"
}

readonly RECORD_USAGE="usage: ,agent-state record [--app APP] [--kind permission|question|auth] <blocked|done|idle|working> [cwd] [message]"

cmd_record() {
    local app="" kind=""
    while [[ ${1:-} == --* ]]; do
        case $1 in
            --app)  app=${2:-}; shift 2 || { echo "$RECORD_USAGE" >&2; return 2; } ;;
            --kind) kind=${2:-}; shift 2 || { echo "$RECORD_USAGE" >&2; return 2; } ;;
            *)      echo "$RECORD_USAGE" >&2; return 2 ;;
        esac
    done

    local state=${1:-} cwd=${2:-$PWD} message=${3:-}

    case $state in
        blocked|done|idle|working) ;;
        *) echo "$RECORD_USAGE" >&2; return 2 ;;
    esac

    # The spec gives kind meaning only on blocked.
    case $kind in
        "") ;;
        permission|question|auth) [[ $state == blocked ]] || kind="" ;;
        *) echo "$RECORD_USAGE" >&2; return 2 ;;
    esac

    [[ $app =~ ^[A-Za-z0-9_.+-]{0,32}$ ]] || { echo "$RECORD_USAGE" >&2; return 2; }

    # The pane is the row identity, so outside tmux there is nothing to record.
    [[ -n "${TMUX_PANE:-}" ]] || return 0

    # Tab and newline are separators in this script's own output, so neither
    # survives into a stored value.
    cwd=$(printf '%s' "$cwd" | tr '\t\n' '  ')
    message=$(printf '%s' "$message" | tr '\t\n' '  ')

    db <<SQL >/dev/null
$PRAGMAS
$SCHEMA
INSERT INTO records (pane, state, kind, app, updated_at, cwd, message)
VALUES ($(sql_quote "$TMUX_PANE"), $(sql_quote "$state"), $(sql_quote "$kind"),
        $(sql_quote "$app"), $(date +%s), $(sql_quote "$cwd"), $(sql_quote "$message"))
ON CONFLICT(pane) DO UPDATE SET
  state      = excluded.state,
  kind       = excluded.kind,
  app        = excluded.app,
  updated_at = excluded.updated_at,
  cwd        = excluded.cwd,
  message    = excluded.message;
SQL
}

cmd_clear() {
    [[ -n "${TMUX_PANE:-}" ]] || return 0
    [[ -f $STATE_DB ]] || return 0

    db <<SQL >/dev/null
$PRAGMAS
$SCHEMA
DELETE FROM records WHERE pane = $(sql_quote "$TMUX_PANE");
DELETE FROM visits WHERE pane = $(sql_quote "$TMUX_PANE");
SQL
}

# Records that a human jumped to this pane. A blocked pane gets a visit, so
# `list` can rank it below the blocked sessions they have not seen yet. A done
# pane becomes idle, since done means the result is unseen and now it has been.
# Called by the two places that switch a tmux client: ,agents and the sketchybar
# agents item.
#
# The pane's state is read here rather than trusted from the caller. Switching
# a client fires no Claude Code event, so the stored state is still current at the
# moment of the jump, and a caller holding a stale row cannot demote a session
# that has since gone back to working.
cmd_visited() {
    local pane=${1:-}

    [[ $pane =~ ^%[0-9]+$ ]] || return 0
    [[ -f $STATE_DB ]] || return 0

    # The WHERE clause is what lets this parse: in an INSERT ... SELECT, SQLite
    # cannot otherwise tell ON CONFLICT from the ON of a join.
    db <<SQL >/dev/null
$PRAGMAS
$SCHEMA
INSERT INTO visits (pane, visited_at)
SELECT pane, $(date +%s) FROM records
WHERE pane = $(sql_quote "$pane") AND state = 'blocked'
ON CONFLICT(pane) DO UPDATE SET visited_at = excluded.visited_at;
-- updated_at stays put, so the age still counts from when the turn finished.
UPDATE records SET state = 'idle'
WHERE pane = $(sql_quote "$pane") AND state = 'done';
SQL
}

# Emits one tab-separated row per live session:
#   state  session:window  dir  age  message  pane  app  kind
#
# Ordered by state (blocked, done, idle, then working) and, within a state,
# oldest first — the session sitting in it longest is the one most likely to
# have been forgotten.
#
# Except for sessions already visited while blocked, which sink below the rest
# of their state, least recently visited first. Without that, jumping to a
# blocked session and leaving it blocked would land on the same row the next
# time the picker opens, since visiting a pane does not change its age.
cmd_list() {
    [[ -f $STATE_DB ]] || return 0

    local live
    live=$(tmux list-panes -a -F '#{pane_id}	#{session_name}:#{window_index}' 2>/dev/null) || return 0
    [[ -n $live ]] || return 0

    # The live pane list goes into a temp table so the join, the prune and the
    # ordering all happen in one statement. A row whose pane is gone belongs to
    # a session that died without firing SessionEnd.
    local values=""
    local pane label
    while IFS=$'\t' read -r pane label; do
        # tmux pane ids are always %<digits>; anything else is not something to
        # interpolate into SQL.
        [[ $pane =~ ^%[0-9]+$ ]] || continue
        values+="($(sql_quote "$pane"), $(sql_quote "$label")),"
    done <<<"$live"
    [[ -n $values ]] || return 0

    db <<SQL | format_rows
$PRAGMAS
$SCHEMA
CREATE TEMP TABLE live (pane TEXT PRIMARY KEY, label TEXT NOT NULL);
INSERT OR IGNORE INTO live (pane, label) VALUES ${values%,};
DELETE FROM records WHERE pane NOT IN (SELECT pane FROM live);
DELETE FROM visits WHERE pane NOT IN (SELECT pane FROM live);
.mode list
.separator "\t"
-- A visit counts only until the session next reports, so visited_at older than
-- updated_at means the pane has moved on and belongs back in the age ordering.
-- That also expires the demotion when a session blocks again: recording a state
-- bumps updated_at past the visit, and a fresh block is worth surfacing.
--
-- Sorting the demoted rows on the visit rather than the age is what makes them
-- rotate. max() picks whichever timestamp ranks the row, so both groups sort
-- with one expression: the visit for a demoted row, the age for every other.
SELECT p.state, l.label, p.cwd, p.updated_at, p.message, p.pane, p.app, p.kind
FROM records p JOIN live l ON l.pane = p.pane
LEFT JOIN visits v ON v.pane = p.pane
ORDER BY CASE p.state WHEN 'blocked' THEN 0 WHEN 'done' THEN 1 WHEN 'idle' THEN 2 ELSE 3 END,
         CASE WHEN COALESCE(v.visited_at, 0) >= p.updated_at THEN 1 ELSE 0 END,
         max(COALESCE(v.visited_at, 0), p.updated_at) ASC;
SQL
}

# Turns the raw SELECT into display columns: cwd to its basename, updated_at to
# a relative age. Done in awk rather than SQL because both are string and clock
# work that SQLite makes awkward, and awk -F'\t' keeps empty fields intact.
format_rows() {
    awk -F'\t' -v now="$(date +%s)" '
        function age(secs) {
            if (secs < 60)   return secs "s"
            if (secs < 3600) return int(secs / 60) "m"
            return int(secs / 3600) "h" int((secs % 3600) / 60) "m"
        }
        {
            n = split($3, parts, "/")
            printf "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", $1, $2, parts[n], age(now - $4), $5, $6, $7, $8
        }
    '
}

main() {
    local command=${1:-}
    shift || true

    command -v sqlite3 >/dev/null 2>&1 || {
        echo ",agent-state: sqlite3 not found" >&2
        return 1
    }

    case $command in
        record)  cmd_record "$@" ;;
        clear)   cmd_clear ;;
        list)    cmd_list ;;
        visited) cmd_visited "$@" ;;
        *)
            echo "usage: ,agent-state <record|clear|list|visited PANE>" >&2
            return 2
            ;;
    esac
}

main "$@"