#!/usr/bin/env bash
# ==============================================================================
# setup-compaction-cron.sh
# Declaratively registers the Multi-Agent Nightly Session Compaction in OpenClaw.
# Schedule: Every night at 03:30 Europe/Berlin (cron: 30 3 * * *)
# ==============================================================================

set -euo pipefail

NAMESPACE="${NAMESPACE:-openclaw}"
POD_NAME="${POD_NAME:-openclaw-0}"
CONTAINER_NAME="${CONTAINER_NAME:-openclaw}"

JOB_NAME="nightly-session-compaction"
CRON_SCHEDULE="30 3 * * *"
TIMEZONE="Europe/Berlin"

echo "=== Registering Multi-Agent Nightly Session Compaction in OpenClaw ($POD_NAME) ==="

# 1. Install or ensure python compaction script is present on PVC
echo "Ensuring /home/node/.openclaw/scripts/compact_sessions.py exists..."
kubectl exec -i -n "$NAMESPACE" "$POD_NAME" -c "$CONTAINER_NAME" -- /bin/sh -c 'cat << "EOF" > /home/node/.openclaw/scripts/compact_sessions.py
#!/usr/bin/env python3
"""
Nightly multi-agent session compaction and maintenance script for OpenClaw.
1. Pre-warms OS page cache for SQLite databases on spinning disks.
2. Prunes abandoned isolated cron sessions (>48h old) to prevent token/state bloat.
3. Compacts all active forum topics and main sessions exceeding the token threshold.
4. Checkpoints SQLite WAL journals across all agent databases to disk.
"""
import glob
import json
import os
import sqlite3
import subprocess
import sys
import time

THRESHOLD_TOKENS = 25000
CRON_MAX_AGE_HOURS = 48
RPC_TIMEOUT_MS = 120000

def prewarm_page_cache():
    print("=== Pre-warming SQLite databases into OS page cache ===")
    t0 = time.time()
    try:
        subprocess.run(
            "cat /home/node/.openclaw/state/openclaw.sqlite /home/node/.openclaw/agents/*/agent/openclaw-agent.sqlite > /dev/null 2>&1 || true",
            shell=True,
            timeout=90.0
        )
        print(f"Page cache pre-warmed in {time.time() - t0:.2f}s")
    except Exception as e:
        print(f"Page cache pre-warming warning: {e}")

def prune_stale_cron_sessions():
    print(f"=== Pruning isolated cron sessions older than {CRON_MAX_AGE_HOURS}h ===")
    cutoff_ms = int(time.time() * 1000) - (CRON_MAX_AGE_HOURS * 3600 * 1000)
    total_pruned = 0
    total_tokens_reclaimed = 0

    for db in sorted(glob.glob("/home/node/.openclaw/agents/*/agent/openclaw-agent.sqlite")):
        agent = db.split("/")[5]
        try:
            conn = sqlite3.connect(db, timeout=30.0)
            conn.execute("PRAGMA foreign_keys = ON;")
            cursor = conn.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type=\"table\"")
            all_tables = set(r[0] for r in cursor.fetchall())

            cursor.execute("SELECT session_key, current_session_id, updated_at, entry_json FROM session_nodes WHERE session_key LIKE \"%:cron:%\"")
            stale = []
            for k, sid, up, ej in cursor.fetchall():
                if up and up < cutoff_ms:
                    tok = 0
                    if ej:
                        try:
                            tok = json.loads(ej).get("totalTokens", 0) or 0
                        except Exception:
                            pass
                    stale.append((k, sid, tok))

            if stale:
                tokens_reclaimed = sum(t[2] for t in stale)
                for k, sid, _ in stale:
                    for t in ["trajectory_runtime_events", "transcript_events", "transcript_event_identities", "session_transcript_active_events", "transcript_rewrite_watermarks", "session_transcript_index_state", "session_transcript_fts", "session_conversations"]:
                        if t in all_tables and sid:
                            cursor.execute(f"DELETE FROM {t} WHERE session_id = ?", (sid,))
                    for t in ["session_participants", "heartbeat_outcomes", "session_progress_cards", "board_widgets", "board_tabs"]:
                        if t in all_tables:
                            cursor.execute(f"DELETE FROM {t} WHERE session_key = ?", (k,))
                    if "memory_entry_origins" in all_tables:
                        cursor.execute("DELETE FROM memory_entry_origins WHERE session_key = ? OR session_id = ?", (k, sid or ""))
                    if "session_windows" in all_tables:
                        cursor.execute("DELETE FROM session_windows WHERE session_key = ? OR session_id = ?", (k, sid or ""))
                    cursor.execute("DELETE FROM session_nodes WHERE session_key = ?", (k,))
                conn.commit()
                total_pruned += len(stale)
                total_tokens_reclaimed += tokens_reclaimed
                print(f"  [{agent}] Pruned {len(stale)} dead cron sessions ({tokens_reclaimed} tokens reclaimed)")

            # Unconditionally resolve any dangling foreign key references across all tables
            while True:
                cursor.execute("PRAGMA foreign_key_check;")
                violations = cursor.fetchall()
                if not violations:
                    break
                for table, rowid, parent, fkid in violations:
                    cursor.execute(f"DELETE FROM {table} WHERE rowid = ?", (rowid,))
            conn.commit()
            conn.close()
        except Exception as e:
            print(f"  [{agent}] Pruning error: {e}")

    # Clear resolved quarantine entries if all databases are healthy
    try:
        q_path = "/home/node/.openclaw/state/openclaw-quarantine.sqlite"
        if os.path.exists(q_path):
            q_conn = sqlite3.connect(q_path)
            q_cur = q_conn.cursor()
            q_cur.execute("DELETE FROM quarantined_databases")
            q_conn.commit()
            q_conn.close()
    except Exception as e:
        print(f"Quarantine cleanup error: {e}")

    print(f"Cron pruning complete: {total_pruned} sessions pruned, {total_tokens_reclaimed} tokens reclaimed.")

def compact_active_sessions():
    print(f"=== Running OpenClaw Multi-Agent Compaction (threshold: {THRESHOLD_TOKENS} tokens, LLM summarization) ===")
    res = subprocess.run(["openclaw", "sessions", "--all-agents", "--json"], capture_output=True, text=True)
    if res.returncode != 0:
        print(f"Failed to list sessions: {res.stderr}")
        return

    sessions = json.loads(res.stdout).get("sessions", [])
    compacted = 0
    skipped = 0

    for s in sessions:
        key = s.get("key", "")
        agent = s.get("agentId") or "main"
        tokens = s.get("totalTokens") or 0

        # Target main sessions and all telegram sessions (topics and direct)
        is_target = (
            key == "agent:main:main" or 
            ":telegram:" in key or
            (key.startswith("agent:") and key.endswith(":main"))
        )
        if is_target and not key.endswith(":heartbeat"):
            if tokens >= THRESHOLD_TOKENS:
                print(f"Compacting {agent} ({tokens} tokens) -> {key}...")
                c_res = subprocess.run(
                    ["openclaw", "sessions", "compact", key, "--agent", agent, "--timeout", str(RPC_TIMEOUT_MS)],
                    capture_output=True, text=True
                )
                if c_res.returncode == 0:
                    print(f"  ✓ {c_res.stdout.strip()}")
                    compacted += 1
                else:
                    print(f"  ✗ {c_res.stderr.strip()}; retrying with --max-lines 200 fallback...")
                    fb_res = subprocess.run(
                        ["openclaw", "sessions", "compact", key, "--agent", agent, "--max-lines", "200", "--timeout", str(RPC_TIMEOUT_MS)],
                        capture_output=True, text=True
                    )
                    if fb_res.returncode == 0:
                        print(f"  ✓ (fallback) {fb_res.stdout.strip()}")
                        compacted += 1
                    else:
                        print(f"  ✗ Fallback failed: {fb_res.stderr.strip()}")
                time.sleep(2)
            else:
                skipped += 1

    print(f"Compaction complete: {compacted} compacted, {skipped} skipped.")

def checkpoint_wal_journals():
    print("=== Checkpointing SQLite WAL journals across all agent databases ===")
    dbs = ["/home/node/.openclaw/state/openclaw.sqlite"] + glob.glob("/home/node/.openclaw/agents/*/agent/openclaw-agent.sqlite")
    wal_freed = 0
    for db in dbs:
        wal = db + "-wal"
        if os.path.exists(wal):
            sz = os.path.getsize(wal)
            try:
                con = sqlite3.connect(db, timeout=30.0)
                con.execute("PRAGMA wal_checkpoint(TRUNCATE);")
                con.close()
                wal_freed += sz
            except Exception as e:
                print(f"  ✗ Error checkpointing {db}: {e}")
    print(f"WAL Checkpoint complete: {wal_freed / (1024*1024):.2f} MB merged to disk.")

def vacuum_and_optimize():
    print("=== Running SQLite VACUUM and OPTIMIZE across all databases ===")
    dbs = ["/home/node/.openclaw/state/openclaw.sqlite"] + glob.glob("/home/node/.openclaw/agents/*/agent/openclaw-agent.sqlite")
    for db in dbs:
        try:
            con = sqlite3.connect(db, timeout=60.0)
            con.execute("PRAGMA optimize;")
            con.execute("VACUUM;")
            con.close()
        except Exception as e:
            print(f"  ✗ Error vacuuming {db}: {e}")
    print("SQLite VACUUM and optimization complete.")

def main():
    prewarm_page_cache()
    try:
        subprocess.run(["python3", "/home/node/.openclaw/scripts/recover_sessions.py"], check=False)
    except Exception:
        pass
    prune_stale_cron_sessions()
    compact_active_sessions()
    checkpoint_wal_journals()
    vacuum_and_optimize()

if __name__ == "__main__":
    main()
EOF
chmod +x /home/node/.openclaw/scripts/compact_sessions.py
'

# 2. Register / Update the OpenClaw cron job
echo "Registering OpenClaw Cron automation '$JOB_NAME' ($CRON_SCHEDULE @ $TIMEZONE)..."
kubectl exec -n "$NAMESPACE" "$POD_NAME" -c "$CONTAINER_NAME" -- openclaw cron add \
  --name "Nightly Session Compaction" \
  --declaration-key "$JOB_NAME" \
  --description "Nightly multi-agent session compaction and SQLite WAL checkpointing" \
  --cron "$CRON_SCHEDULE" \
  --tz "$TIMEZONE" \
  --session "isolated" \
  --no-deliver \
  --command "python3 /home/node/.openclaw/scripts/compact_sessions.py"

echo "=== Compaction Cron automation successfully registered! ==="
kubectl exec -n "$NAMESPACE" "$POD_NAME" -c "$CONTAINER_NAME" -- openclaw cron show "$JOB_NAME" 2>/dev/null || true