#!/usr/bin/env bash
# Diagnose compact context state for the current workspace.
# Usage: bash check.sh [workspace_dir]

WORKSPACE="${1:-$(pwd)}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SET_KEEP_SCRIPT="$SCRIPT_DIR/set-keep-recent.mjs"
SETTINGS_FILE="$WORKSPACE/.pi/settings.json"
AGENTS_MD="$WORKSPACE/AGENTS.md"
PERSONAL_AGENTS="$HOME/.agents/AGENTS.md"
OPENCLAW_CFG="${OPENCLAW_CONFIG_PATH:-$HOME/.openclaw/openclaw.json}"
STATE_DIR="${OPENCLAW_STATE_DIR:-$HOME/.openclaw}"
BACKUP_KEY="__compactContextSkillBackup"
LAST_RUN_KEY="__compactContextSkillLastRun"
LAST_RESTORE_KEY="__compactContextSkillLastRestore"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
RESET='\033[0m'

echo ""
echo -e "${BOLD}━━━ Compact Context Diagnostics ━━━${RESET}"
echo ""

# ── 0. Current session selector ───────────────────────────────────────────────
echo -e "${BOLD}Current session selector:${RESET}"
if [ -n "${OPENCLAW_TASK_SESSION_KEY:-}" ]; then
  echo "  session_key=$OPENCLAW_TASK_SESSION_KEY"
elif [ -n "${OPENCLAW_TASK_SESSION_ID:-}" ]; then
  echo "  session_id=$OPENCLAW_TASK_SESSION_ID"
else
  echo -e "  ${YELLOW}○ No OPENCLAW_TASK_SESSION_KEY / OPENCLAW_TASK_SESSION_ID in this shell${RESET}"
fi

# ── 1. Compact history from agent session stores ──────────────────────────────
echo -e "${BOLD}Compact history:${RESET}"
if [ -d "$STATE_DIR/agents" ]; then
  python3 - "$STATE_DIR" <<'PY'
import json
import os
import sys

state_dir = sys.argv[1]
total = 0
sessions = 0
stores = 0

agents_dir = os.path.join(state_dir, "agents")
for agent_id in os.listdir(agents_dir):
    store = os.path.join(agents_dir, agent_id, "sessions", "sessions.json")
    if not os.path.exists(store):
        continue
    stores += 1
    try:
        with open(store, "r", encoding="utf-8") as handle:
            data = json.load(handle)
    except Exception as exc:
        print(f"  unable to parse {store}: {exc}")
        continue
    for key, value in data.items():
        if key in {"global", "unknown"} or not isinstance(value, dict):
            continue
        count = value.get("compactionCount", 0) or 0
        total += count
        if count > 0:
            sessions += 1

print(f"  stores={stores} total_compactions={total} sessions_with_compactions={sessions}")
PY
else
  echo -e "  ${YELLOW}○ No agent session stores found under${RESET} $STATE_DIR"
fi

# ── 2. OpenClaw sessions CLI availability ─────────────────────────────────────
echo ""
echo -e "${BOLD}Native session CLI:${RESET}"
if command -v openclaw >/dev/null 2>&1; then
  CLI_OUTPUT="$(openclaw sessions --json 2>&1 || true)"
  CLI_STATUS=$(CLI_OUTPUT="$CLI_OUTPUT" python3 - <<'PY'
import json
import os

text = os.environ.get("CLI_OUTPUT", "")
payload = None
for start, start_char in enumerate(text):
    if start_char not in "{[":
        continue
    stack = []
    in_string = False
    escaped = False
    json_text = None
    for index in range(start, len(text)):
        char = text[index]
        if in_string:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
            continue
        if char == '"':
            in_string = True
            continue
        if char in "{[":
            stack.append("}" if char == "{" else "]")
            continue
        if char in "}]":
            if not stack or stack.pop() != char:
                json_text = None
                break
            if not stack:
                json_text = text[start:index + 1]
                break
    if json_text is None:
        continue
    try:
        candidate = json.loads(json_text)
    except Exception:
        continue
    if isinstance(candidate, dict) and ("sessions" in candidate or "path" in candidate):
        payload = candidate
        break

if payload is None:
    print("unavailable")
    raise SystemExit(0)

sessions = payload.get("sessions") or []
path = payload.get("path") or "unknown"
print(f"available path={path} sessions={len(sessions)}")
PY
)
  if [[ "$CLI_STATUS" == available* ]]; then
    echo -e "  ${GREEN}✔${RESET} $CLI_STATUS"
  else
    echo -e "  ${RED}✘${RESET} openclaw sessions --json unavailable"
  fi
else
  echo -e "  ${RED}✘${RESET} openclaw command not found"
fi

# ── 3. Session alignment / token accounting ───────────────────────────────────
echo ""
echo -e "${BOLD}Current session alignment:${RESET}"
if [ -f "$SET_KEEP_SCRIPT" ]; then
  INFO_OUTPUT="$(node "$SET_KEEP_SCRIPT" --info --workspace "$WORKSPACE" --allow-latest-updated --debug 2>&1 || true)"
  if [[ -n "$INFO_OUTPUT" ]]; then
    INFO_REPORT=$(INFO_OUTPUT="$INFO_OUTPUT" python3 - <<'PY'
import os
import re

text = os.environ.get("INFO_OUTPUT", "")
pairs = dict(re.findall(r'([A-Za-z_]+)=([^\s]+)', text))

def num(name):
    value = pairs.get(name)
    if value in (None, "unknown", "none"):
        return None
    try:
        return float(value)
    except ValueError:
        return None

native = num("native_total_tokens")
estimate = num("transcript_total_estimate")
bias = None
ratio = None
if native is not None and estimate is not None and native != 0:
    bias = estimate - native
    ratio = estimate / native

lines = [
    f"source={pairs.get('source', 'unknown')}",
    f"estimation_mode={pairs.get('estimation_mode', 'unknown')}",
    f"source_label={pairs.get('source_label', 'unknown')}",
    f"display_label={pairs.get('display_label', 'unknown')}",
    f"prompt_truth_source={pairs.get('prompt_truth_source', 'unknown')}",
    f"resolved_by={pairs.get('resolved_by', 'unknown')}",
    f"session_key={pairs.get('resolved_session_key', 'unknown')}",
    f"session_id={pairs.get('resolved_session_id', 'unknown')}",
    f"session_file={pairs.get('resolved_session_file', 'unknown')}",
    f"turns={pairs.get('turns', 'unknown')}",
    f"total_tokens={pairs.get('total_tokens', 'unknown')}",
    f"message_estimate_before={pairs.get('message_estimate_before', 'unknown')}",
    f"message_estimate_after={pairs.get('message_estimate_after', 'unknown')}",
    f"context_tokens={pairs.get('context_tokens', 'unknown')}",
    f"used_pct={pairs.get('used_pct', 'unknown')}",
    f"snapshot_status={pairs.get('snapshot_status', 'unknown')}",
    f"native_total_tokens={pairs.get('native_total_tokens', 'unknown')}",
    f"transcript_total_estimate={pairs.get('transcript_total_estimate', 'unknown')}",
]

if bias is not None and ratio is not None:
    lines.append(f"estimate_bias={bias:.0f}")
    lines.append(f"estimate_ratio={ratio:.3f}")

print("\n".join(lines))
PY
)
    while IFS= read -r line; do
      [ -n "$line" ] && echo "  $line"
    done <<< "$INFO_REPORT"
  else
    echo -e "  ${RED}✘${RESET} set-keep-recent --info produced no output"
  fi
else
  echo -e "  ${RED}✘${RESET} set-keep-recent.mjs not found"
fi

# ── 4. Backup / workspace settings ────────────────────────────────────────────
echo ""
echo -e "${BOLD}Workspace settings:${RESET} $SETTINGS_FILE"
if [ -f "$SETTINGS_FILE" ]; then
  python3 - "$SETTINGS_FILE" <<'PY'
import json
import sys

settings_file = sys.argv[1]
with open(settings_file, "r", encoding="utf-8") as handle:
    data = json.load(handle)

compaction = data.get("compaction") or {}
keep_recent = compaction.get("keepRecentTokens", "not set")
backup = compaction.get("__compactContextSkillBackup")
last_run = compaction.get("__compactContextSkillLastRun")
last_restore = compaction.get("__compactContextSkillLastRestore")

print(f"  keepRecentTokens={keep_recent}")
if isinstance(backup, dict):
    print("  backup=present")
    print(f"  backup.sessionKey={backup.get('sessionKey', 'unknown')}")
    print(f"  backup.sessionId={backup.get('sessionId', 'unknown')}")
    print(f"  backup.source={backup.get('source', 'unknown')}")
    print(f"  backup.appliedKeepRecentTokens={backup.get('appliedKeepRecentTokens', 'unknown')}")
else:
    print("  backup=absent")

if isinstance(last_run, dict):
    print("  last_run=present")
    print(f"  last_run.source={last_run.get('source', 'unknown')}")
    print(f"  last_run.sessionKey={last_run.get('sessionKey', 'unknown')}")
    print(f"  last_run.totalTokensBefore={last_run.get('totalTokensBefore', 'unknown')}")
    print(f"  last_run.usedPctBefore={last_run.get('usedPctBefore', 'unknown')}")
    print(f"  last_run.keepRecentTokens={last_run.get('keepRecentTokens', 'unknown')}")
    print(f"  last_run.messageEstimateBefore={last_run.get('messageEstimateBefore', 'unknown')}")
    print(f"  last_run.messageEstimateAfter={last_run.get('messageEstimateAfter', 'unknown')}")
    print(f"  last_run.promptTruthSource={last_run.get('promptTruthSource', 'unknown')}")
else:
    print("  last_run=absent")

if isinstance(last_restore, dict):
    print("  last_restore=present")
    print(f"  last_restore.restoreStatus={last_restore.get('restoreStatus', 'unknown')}")
    print(f"  last_restore.restoredKeepRecentTokens={last_restore.get('restoredKeepRecentTokens', 'unknown')}")
else:
    print("  last_restore=absent")
PY
else
  echo -e "  ${YELLOW}○ settings.json not found${RESET}"
fi

# ── 5. Workspace / personal AGENTS.md ─────────────────────────────────────────
echo ""
echo -e "${BOLD}Workspace AGENTS.md:${RESET} $AGENTS_MD"
if [ -f "$AGENTS_MD" ]; then
  if grep -q "## Session Startup" "$AGENTS_MD" 2>/dev/null; then
    if grep -q "Start your reply with:" "$AGENTS_MD" 2>/dev/null || grep -q "（已压缩，从摘要恢复）" "$AGENTS_MD" 2>/dev/null; then
      echo -e "  ${YELLOW}○ ## Session Startup${RESET} — legacy (fixed-prefix rule detected)"
    else
      echo -e "  ${GREEN}✔ ## Session Startup${RESET} — safe"
    fi
  else
    echo -e "  ${RED}✘ ## Session Startup${RESET} — missing"
  fi
else
  echo -e "  ${YELLOW}○ AGENTS.md not found in workspace${RESET}"
fi

echo ""
echo -e "${BOLD}Personal AGENTS.md:${RESET} $PERSONAL_AGENTS"
if [ -f "$PERSONAL_AGENTS" ]; then
  if grep -q "## Session Startup" "$PERSONAL_AGENTS" 2>/dev/null; then
    if grep -q "Start your reply with:" "$PERSONAL_AGENTS" 2>/dev/null || grep -q "（已压缩，从摘要恢复）" "$PERSONAL_AGENTS" 2>/dev/null; then
      echo -e "  ${YELLOW}○ ## Session Startup${RESET} — legacy (fixed-prefix rule detected)"
    else
      echo -e "  ${GREEN}✔ ## Session Startup${RESET} — safe"
    fi
  else
    echo -e "  ${RED}✘ ## Session Startup${RESET} — missing"
  fi
else
  echo -e "  ${YELLOW}○ personal AGENTS.md not found${RESET}"
fi

# ── 6. memoryFlush config ─────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}memoryFlush config:${RESET}"
if [ -f "$OPENCLAW_CFG" ]; then
  python3 - "$OPENCLAW_CFG" <<'PY'
import json
import sys

cfg_path = sys.argv[1]
with open(cfg_path, "r", encoding="utf-8") as handle:
    cfg = json.load(handle)

memory_flush = cfg.get("agents", {}).get("defaults", {}).get("compaction", {}).get("memoryFlush", {})
if memory_flush:
    print(f"  enabled={memory_flush.get('enabled', True)} softThresholdTokens={memory_flush.get('softThresholdTokens', 'not set')}")
else:
    print("  not configured")
PY
else
  echo -e "  ${YELLOW}○ openclaw.json not found${RESET}"
fi

# ── 7. Summary ────────────────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}━━━ Summary ━━━${RESET}"

ISSUES=0

if ! command -v openclaw >/dev/null 2>&1; then
  echo -e "${RED}✘ Native OpenClaw CLI unavailable.${RESET}"
  ISSUES=$((ISSUES + 1))
fi

if [ -z "${OPENCLAW_TASK_SESSION_KEY:-}" ] && [ -z "${OPENCLAW_TASK_SESSION_ID:-}" ]; then
  echo -e "${YELLOW}○ Current shell is missing a bound session selector.${RESET}"
  echo -e "  ${YELLOW}→ Strict skill mode will refuse to run here until the live thread exports it.${RESET}"
  ISSUES=$((ISSUES + 1))
fi

if [ -f "$SETTINGS_FILE" ] && grep -q "\"$BACKUP_KEY\"" "$SETTINGS_FILE" 2>/dev/null; then
  echo -e "${CYAN}○${RESET} Workspace currently contains a compact-context backup entry."
fi

if [ -f "$SETTINGS_FILE" ] && grep -q "\"$LAST_RUN_KEY\"" "$SETTINGS_FILE" 2>/dev/null; then
  echo -e "${CYAN}○${RESET} Workspace contains a last-run snapshot for compaction diagnostics."
fi

if [ -f "$AGENTS_MD" ] && grep -q "## Session Startup" "$AGENTS_MD" 2>/dev/null; then
  STARTUP_OK=1
else
  STARTUP_OK=0
fi

if [ -f "$PERSONAL_AGENTS" ] && grep -q "## Session Startup" "$PERSONAL_AGENTS" 2>/dev/null; then
  PERSONAL_OK=1
else
  PERSONAL_OK=0
fi

if ([ -f "$AGENTS_MD" ] && (grep -q "Start your reply with:" "$AGENTS_MD" 2>/dev/null || grep -q "（已压缩，从摘要恢复）" "$AGENTS_MD" 2>/dev/null)) \
  || ([ -f "$PERSONAL_AGENTS" ] && (grep -q "Start your reply with:" "$PERSONAL_AGENTS" 2>/dev/null || grep -q "（已压缩，从摘要恢复）" "$PERSONAL_AGENTS" 2>/dev/null)); then
  echo -e "${YELLOW}○ Legacy startup rules detected.${RESET}"
  echo -e "  ${YELLOW}→ Run: bash setup.sh${RESET}"
fi

if [ "$STARTUP_OK" -gt 0 ] || [ "$PERSONAL_OK" -gt 0 ]; then
  echo -e "${GREEN}✔ Post-compaction re-injection is set up.${RESET}"
else
  echo -e "${RED}✘ Post-compaction re-injection is NOT set up.${RESET}"
  echo -e "  ${YELLOW}→ Run: bash setup.sh${RESET}"
  ISSUES=$((ISSUES + 1))
fi

echo -e "${CYAN}○${RESET} Token accounting note: skill messageEstimate* / keepRecentTokens are display-side compaction estimates;"
echo -e "  ${CYAN}sessions.json totalTokens${RESET} reflects runtime prompt/cache usage and should be treated as authoritative only after the next real model call."

if printf '%s' "$INFO_OUTPUT" | grep -q "source=transcript-estimate" 2>/dev/null; then
  echo -e "${YELLOW}○ Current token accounting is using transcript-estimate fallback.${RESET}"
fi

echo ""
[ "$ISSUES" -eq 0 ] && echo -e "${GREEN}All checks passed.${RESET}" || echo -e "${YELLOW}$ISSUES issue(s) found — see above.${RESET}"
echo ""
