#!/usr/bin/env bash
# AY Framework -- UserPromptSubmit hook: context-pressure warning (task 31)
#
# Estimates how full the context window is and warns the human BEFORE the next
# turn so they can /compact while there is still headroom:
#   >= 60%  yellow  "⚠ Context at X% — consider /compact before the next phase"
#   >= 85%  red     "🔴 Context critical — run /compact now or risk losing work"
#
# ADVISORY ONLY: prints to stderr and always exits 0. It never blocks a prompt and
# never writes to stdout (stdout on UserPromptSubmit is injected into the model's
# context — the opposite of what we want when context is already tight). If it can't
# determine usage, it stays silent.
#
# Token source (first that works wins):
#   1. $CLAUDE_CONTEXT_TOKENS               — explicit override / testing
#   2. .transcript_path from the hook's stdin JSON  — the live session transcript
#   3. ~/.claude/projects/**/$CLAUDE_CODE_SESSION_ID.jsonl  — fallback lookup
# The window size defaults to 200000 and can be overridden with
# $CLAUDE_CONTEXT_WINDOW (set it to 1000000 on a 1M-context model).
#
# Works on Mac, Linux, Windows (Git Bash). No non-portable tools (no `tac`).

set -uo pipefail

WINDOW="${CLAUDE_CONTEXT_WINDOW:-${CLAUDE_CONTEXT_LIMIT:-200000}}"
[ "$WINDOW" -gt 0 ] 2>/dev/null || WINDOW=200000

# --- read the hook event JSON from stdin (UserPromptSubmit passes transcript_path)
INPUT="$(cat 2>/dev/null || true)"

# --- helper: pull a string field out of the stdin JSON (jq, then python3, then sed)
json_field() {
  local key="$1"
  if [ -z "$INPUT" ]; then return 1; fi
  if command -v jq >/dev/null 2>&1; then
    printf '%s' "$INPUT" | jq -r --arg k "$key" '.[$k] // empty' 2>/dev/null && return 0
  fi
  if command -v python3 >/dev/null 2>&1; then
    printf '%s' "$INPUT" | python3 -c 'import sys,json
try:
    print(json.load(sys.stdin).get(sys.argv[1],"") or "")
except Exception:
    pass' "$key" 2>/dev/null && return 0
  fi
  printf '%s' "$INPUT" | sed -nE "s/.*\"$key\"[[:space:]]*:[[:space:]]*\"([^\"]*)\".*/\1/p" | head -1
}

# --- helper: sum the last usage entry (input + cache_creation + cache_read) --------
usage_total() {
  local f="$1"
  [ -f "$f" ] || return 1
  if command -v jq >/dev/null 2>&1; then
    jq -rn '[inputs
             | select(.message.usage != null) | .message.usage
             | (.input_tokens // 0) + (.cache_creation_input_tokens // 0) + (.cache_read_input_tokens // 0)]
            | last // empty' "$f" 2>/dev/null
    return 0
  fi
  # Pure-grep fallback: take the last usage object and add the three token counts.
  local last it cc cr
  last="$(grep -o '"usage":{[^{]*"cache_read_input_tokens":[0-9]*' "$f" 2>/dev/null | tail -1)"
  [ -n "$last" ] || return 1
  it="$(printf '%s' "$last" | grep -oE '"input_tokens":[0-9]+' | grep -oE '[0-9]+' | head -1)"
  cc="$(printf '%s' "$last" | grep -oE '"cache_creation_input_tokens":[0-9]+' | grep -oE '[0-9]+' | head -1)"
  cr="$(printf '%s' "$last" | grep -oE '"cache_read_input_tokens":[0-9]+' | grep -oE '[0-9]+' | head -1)"
  echo $(( ${it:-0} + ${cc:-0} + ${cr:-0} ))
}

# --- resolve the current token count ----------------------------------------------
TOKENS=""
if [ -n "${CLAUDE_CONTEXT_TOKENS:-}" ]; then
  TOKENS="$CLAUDE_CONTEXT_TOKENS"
else
  TRANSCRIPT="$(json_field transcript_path)"
  if [ -z "$TRANSCRIPT" ] || [ ! -f "$TRANSCRIPT" ]; then
    # Fallback: locate the transcript by session id under ~/.claude/projects/
    SID="${CLAUDE_CODE_SESSION_ID:-}"
    if [ -n "$SID" ] && [ -d "$HOME/.claude/projects" ]; then
      TRANSCRIPT="$(find "$HOME/.claude/projects" -name "${SID}.jsonl" 2>/dev/null | head -1)"
    fi
  fi
  [ -n "$TRANSCRIPT" ] && TOKENS="$(usage_total "$TRANSCRIPT" 2>/dev/null || true)"
fi

# Can't determine usage → stay silent, never block.
[ -n "$TOKENS" ] || exit 0
case "$TOKENS" in ''|*[!0-9]*) exit 0 ;; esac

# --- compute percentage -----------------------------------------------------------
PCT=$(( TOKENS * 100 / WINDOW ))

# --- locate .ay/ (walk up from cwd) so we can read config + write the snapshot ----
find_ay_dir() {
  local dir="$PWD"
  while [ -n "$dir" ] && [ "$dir" != "/" ]; do
    [ -d "$dir/.ay" ] && { printf '%s\n' "$dir/.ay"; return 0; }
    dir="$(dirname "$dir")"
  done
  [ -d "${HOME:-}/.ay" ] && { printf '%s\n' "$HOME/.ay"; return 0; }
  return 1
}
AY_DIR="$(find_ay_dir || true)"

# --- pause threshold: .ay/config.json key context_pressure_threshold (default 0.8) -
# Accepts a fraction (0.8) or a whole percent (80); normalizes to an integer percent.
PAUSE_PCT=80
if [ -n "$AY_DIR" ] && [ -f "$AY_DIR/config.json" ]; then
  RAW=""
  if command -v jq >/dev/null 2>&1; then
    RAW="$(jq -r '.context_pressure_threshold // empty' "$AY_DIR/config.json" 2>/dev/null || true)"
  fi
  [ -n "$RAW" ] || RAW="$(grep -oE '"context_pressure_threshold"[[:space:]]*:[[:space:]]*[0-9.]+' "$AY_DIR/config.json" 2>/dev/null | grep -oE '[0-9.]+$' | head -1)"
  if [ -n "$RAW" ]; then
    PAUSE_PCT="$(awk -v v="$RAW" 'BEGIN{ n=v+0; if (n<=1) n=n*100; printf "%d", n+0.5 }')"
  fi
fi
[ "$PAUSE_PCT" -gt 0 ] 2>/dev/null || PAUSE_PCT=80

# --- auto-pause snapshot (task 65) -------------------------------------------------
# When pressure crosses the pause threshold, write a resumable state snapshot to
# .ay/tracking/context-snapshot-{ts}.json — once per session (a per-session marker
# guards against re-writing on every prompt). Sets SNAP_PATH for the caller.
SNAP_PATH=""
write_snapshot() {
  SNAP_PATH=""
  [ -n "$AY_DIR" ] || return 0
  local tdir="$AY_DIR/tracking"
  mkdir -p "$tdir" 2>/dev/null || true

  # Once per pause-episode: marker keyed by session id (fallback: shared marker).
  local sid="${CLAUDE_CODE_SESSION_ID:-nosid}"
  local marker="$tdir/.context-pause-${sid}"
  if [ -f "$marker" ]; then
    SNAP_PATH="$(cat "$marker" 2>/dev/null || true)"
    return 0
  fi

  local ts_file ts_iso task phase branch
  ts_file="$(date +%Y%m%dT%H%M%S 2>/dev/null || echo now)"
  ts_iso="$(date -u +%Y-%m-%dT%H:%M:%SZ 2>/dev/null || echo '')"
  SNAP_PATH="$tdir/context-snapshot-${ts_file}.json"

  # current task id from the active lock (tracking/locks preferred, legacy fallback).
  # In a parallel swarm several locks may be present; prefer the one owned by THIS
  # session (matching CLAUDE_CODE_SESSION_ID against the lock's session:/id field),
  # and only fall back to the first lock when ownership can't be established.
  task=""; local lockpath=""
  local d f first_task="" first_path="" sid_lock="${CLAUDE_CODE_SESSION_ID:-}"
  for d in "$AY_DIR/tracking/locks" "$AY_DIR/locks"; do
    [ -d "$d" ] || continue
    for f in "$d"/*.lock; do
      [ -f "$f" ] || continue
      [ -n "$first_task" ] || { first_task="$(basename "$f" .lock)"; first_path="$f"; }
      if [ -n "$sid_lock" ] && grep -qiE "(session|id):[[:space:]]*${sid_lock}([[:space:]]|$)" "$f" 2>/dev/null; then
        task="$(basename "$f" .lock)"; lockpath="$f"; break
      fi
    done
    [ -n "$task" ] && break
  done
  [ -n "$task" ] || { task="$first_task"; lockpath="$first_path"; }

  # last completed phase (best-effort: active lock's phase: line, then state.json)
  phase=""
  if [ -n "$lockpath" ] && [ -f "$lockpath" ]; then
    phase="$(grep -ioE 'phase:[[:space:]]*[A-Za-z0-9_-]+' "$lockpath" 2>/dev/null | head -1 | sed -E 's/^[Pp]hase:[[:space:]]*//')"
  fi
  if [ -z "$phase" ] && [ -f "$AY_DIR/state.json" ]; then
    phase="$(grep -oE '"(last_completed_phase|current_phase)"[[:space:]]*:[[:space:]]*"[^"]*"' "$AY_DIR/state.json" 2>/dev/null | head -1 | sed -E 's/.*"([^"]*)"$/\1/')"
  fi

  # staged files (JSON array), branch
  branch="$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo '')"
  local staged_json="[]"
  if git rev-parse --show-toplevel >/dev/null 2>&1; then
    if command -v jq >/dev/null 2>&1; then
      staged_json="$(git diff --cached --name-only 2>/dev/null | jq -R . | jq -sc . 2>/dev/null || echo '[]')"
      [ -n "$staged_json" ] || staged_json="[]"
    else
      local first=1 line out="["
      while IFS= read -r line; do
        [ -n "$line" ] || continue
        line="${line//\\/\\\\}"; line="${line//\"/\\\"}"
        if [ "$first" -eq 1 ]; then out="$out\"$line\""; first=0; else out="$out,\"$line\""; fi
      done < <(git diff --cached --name-only 2>/dev/null)
      out="$out]"; staged_json="$out"
    fi
  fi

  {
    printf '{\n'
    printf '  "timestamp": "%s",\n' "$ts_iso"
    printf '  "context_pct": %d,\n' "$PCT"
    printf '  "threshold_pct": %d,\n' "$PAUSE_PCT"
    printf '  "tokens": %d,\n' "$TOKENS"
    printf '  "window": %d,\n' "$WINDOW"
    printf '  "task_id": "%s",\n' "$task"
    printf '  "last_completed_phase": "%s",\n' "$phase"
    printf '  "branch": "%s",\n' "$branch"
    printf '  "staged_files": %s\n' "$staged_json"
    printf '}\n'
  } > "$SNAP_PATH" 2>/dev/null || { SNAP_PATH=""; return 0; }

  printf '%s\n' "$SNAP_PATH" > "$marker" 2>/dev/null || true
  return 0
}

# Color only when stderr is a terminal; otherwise emit plain text.
if [ -t 2 ]; then
  YELLOW=$'\033[33m'; RED=$'\033[31m'; RESET=$'\033[0m'
else
  YELLOW=''; RED=''; RESET=''
fi

if [ "$PCT" -ge "$PAUSE_PCT" ]; then
  # Auto-pause: snapshot state, then tell the human how to resume. Pausing supersedes
  # the advisory warnings below (they only fire when the threshold is set above 85%).
  write_snapshot
  printf '%s🛑 PAUSED — context at %d%% (>= %d%% threshold, %d/%d tokens)%s\n' \
    "$RED" "$PCT" "$PAUSE_PCT" "$TOKENS" "$WINDOW" "$RESET" >&2
  if [ -n "$SNAP_PATH" ]; then
    printf '   State snapshot: %s\n' "$SNAP_PATH" >&2
    printf '   Resume: open a fresh session and run /continue (state is preserved), or /compact to keep going.\n' >&2
  else
    printf '   Run /compact now or /pause to snapshot manually before you lose work.\n' >&2
  fi
elif [ "$PCT" -ge 85 ]; then
  printf '%s🔴 Context critical (%d%% — %d/%d tokens) — run /compact now or risk losing work%s\n' \
    "$RED" "$PCT" "$TOKENS" "$WINDOW" "$RESET" >&2
elif [ "$PCT" -ge 60 ]; then
  printf '%s⚠ Context at %d%% (%d/%d tokens) — consider /compact before the next phase%s\n' \
    "$YELLOW" "$PCT" "$TOKENS" "$WINDOW" "$RESET" >&2
fi

exit 0