#!/usr/bin/env python3
"""
Nightly multi-agent session maintenance and compaction script for OpenClaw.
1. Pre-warms OS page cache for SQLite databases on spinning disks.
2. Recovers any wedged or locked sessions via recover_sessions.py.
3. Prunes abandoned isolated cron sessions (>48h old) to prevent token/state bloat.
4. Tier A Hard-Reset: Resets koch, hebamme, and shopping topics to 0 tokens (clean archive).
5. Tier B Smart-Compaction: Compacts active sessions exceeding token thresholds (5s pacing).
6. Checkpoints SQLite WAL journals across all agent databases to disk.
7. Optimizes SQLite query planner without blocking full VACUUM.
"""
import glob
import json
import os
import sqlite3
import subprocess
import sys
import time

THRESHOLD_TOKENS = 15000
THRESHOLD_TRANSCRIPT_BYTES = 200000
THRESHOLD_EVENT_COUNT = 30
CRON_MAX_AGE_HOURS = 48
RPC_TIMEOUT_MS = 180000

# Tracking Channels: Deterministic sliding-window (mila_tracker, nutrition_tracker).
# Strict sliding window of ~5-8 events (2-3 turns max), 0 LLM tokens, instant truncation.
TRACKING_TARGETS = [
    ("mila_tracker", "agent:mila_tracker:telegram:group:-1004326921510:topic:2796"),
    ("nutrition_tracker", "agent:nutrition_tracker:telegram:group:-1004326921510:topic:2797"),
]
MAX_TRACKING_EVENTS = 8
MAX_TRACKING_LINES = 8

# Tier A: High-frequency interaction topics (koch, hebamme, shopping).
# Deep-compaction keeps a compact summary + latest 2 turns without completely wiping conversational memory.
TIER_A_TARGETS = [
    ("koch", "agent:koch:telegram:group:-1004326921510:topic:21"),
    ("hebamme", "agent:hebamme:telegram:group:-1004326921510:topic:22"),
    ("shopping", "agent:shopping:telegram:group:-1004326921510:topic:115"),
]

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_tracking_sessions():
    print("=== Tracking Channels Sliding-Window Maintenance (mila_tracker, nutrition_tracker) ===")
    for agent, key in TRACKING_TARGETS:
        db_path = f"/home/node/.openclaw/agents/{agent}/agent/openclaw-agent.sqlite"
        if not os.path.exists(db_path):
            continue
        try:
            con = sqlite3.connect(db_path, timeout=5.0)
            row = con.execute("SELECT current_session_id FROM session_nodes WHERE session_key=?", (key,)).fetchone()
            if not row or not row[0]:
                con.close()
                print(f"  [{agent}] No active session node found for {key}.")
                continue
            sid = row[0]
            cnt = con.execute("SELECT count(*) FROM transcript_events WHERE session_id=?", (sid,)).fetchone()[0]
            con.close()

            if cnt > MAX_TRACKING_EVENTS:
                print(f"  [{agent}] Session {key} has {cnt} events (> {MAX_TRACKING_EVENTS}). Truncating to last {MAX_TRACKING_LINES} lines...")
                cmd = [
                    "openclaw", "sessions", "compact", key,
                    "--agent", agent,
                    "--max-lines", str(MAX_TRACKING_LINES),
                    "--timeout", "30000"
                ]
                r = subprocess.run(cmd, capture_output=True, text=True)
                if r.returncode == 0:
                    print(f"    ✓ Truncated successfully: {r.stdout.strip()}")
                else:
                    print(f"    ✗ Truncation failed: {r.stderr.strip()}")
            else:
                print(f"  [{agent}] Session {key} has {cnt} events (<= {MAX_TRACKING_EVENTS}). Clean.")
        except Exception as e:
            print(f"  ✗ [{agent}] Error during tracking compaction: {e}")

def compact_tier_a_sessions():
    print("=== Tier A: Deep-Compaction (koch, hebamme, shopping) ===")
    for agent, key in TIER_A_TARGETS:
        db_path = f"/home/node/.openclaw/agents/{agent}/agent/openclaw-agent.sqlite"
        if not os.path.exists(db_path):
            continue
        try:
            con = sqlite3.connect(db_path, timeout=10.0)
            row = con.execute("SELECT current_session_id FROM session_nodes WHERE session_key=?", (key,)).fetchone()
            con.close()
            if not row or not row[0]:
                print(f"  [Tier A] {agent} ({key}) already at 0 (no active session).")
                continue
            ev_cnt, ev_bytes = get_session_metrics(agent, key, row[0])
            if ev_cnt <= 2 and ev_bytes < 20000:
                print(f"  [Tier A] {agent} ({key}) has only {ev_cnt} events ({ev_bytes} bytes), skipping.")
                continue
            print(f"  [Tier A] Deep-compacting {agent} ({key}) [{ev_cnt} events, {ev_bytes//1024}KB]...")
            r = subprocess.run(
                ["openclaw", "sessions", "compact", key, "--agent", agent, "--timeout", str(RPC_TIMEOUT_MS)],
                capture_output=True, text=True
            )
            if r.returncode == 0:
                print(f"    ✓ Deep-compacted: {r.stdout.strip()}")
            elif "Already compacted" in r.stderr or "Already compacted" in r.stdout:
                print(f"    ✓ Already compacted.")
            else:
                print(f"    ✗ Standard compact error: {r.stderr.strip()}; retrying with --max-lines 40...")
                fb = subprocess.run(
                    ["openclaw", "sessions", "compact", key, "--agent", agent, "--max-lines", "40", "--timeout", str(RPC_TIMEOUT_MS)],
                    capture_output=True, text=True
                )
                if fb.returncode == 0:
                    print(f"    ✓ Truncation fallback success: {fb.stdout.strip()}")
                else:
                    print(f"    ✗ Truncation fallback failed: {fb.stderr.strip()}")
        except Exception as e:
            print(f"    ✗ Exception deep-compacting {agent}: {e}")
        time.sleep(2)

def get_session_metrics(agent, key, sid=None):
    db_path = f"/home/node/.openclaw/agents/{agent}/agent/openclaw-agent.sqlite"
    if not os.path.exists(db_path):
        return 0, 0
    try:
        con = sqlite3.connect(db_path, timeout=5.0)
        cur = con.cursor()
        if not sid:
            row = cur.execute("SELECT current_session_id FROM session_nodes WHERE session_key=?", (key,)).fetchone()
            sid = row[0] if row else None
        if not sid:
            con.close()
            return 0, 0
        ev_cnt, ev_bytes = cur.execute(
            "SELECT count(*), coalesce(sum(length(event_json)), 0) FROM transcript_events WHERE session_id=?", 
            (sid,)
        ).fetchone()
        con.close()
        return ev_cnt, ev_bytes
    except Exception:
        return 0, 0

def compact_active_sessions():
    excluded_keys = set(t[1] for t in TIER_A_TARGETS) | set(t[1] for t in TRACKING_TARGETS)
    print(f"=== Tier B: Running Multi-Agent Compaction (threshold: {THRESHOLD_TOKENS} tokens / {THRESHOLD_TRANSCRIPT_BYTES//1024}KB, 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
        sid = s.get("sessionId")

        # Skip Tier A and Tracking agents as they have dedicated routines
        if key in excluded_keys:
            continue

        # 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"):
            ev_cnt, ev_bytes = get_session_metrics(agent, key, sid)
            needs_compaction = (
                tokens >= THRESHOLD_TOKENS or
                ev_bytes >= THRESHOLD_TRANSCRIPT_BYTES or
                (tokens == 0 and ev_cnt >= THRESHOLD_EVENT_COUNT)
            )
            if needs_compaction:
                stat_str = f"{tokens} tokens" if tokens else f"{ev_bytes//1024}KB / {ev_cnt} events"
                print(f"Compacting {agent} ({stat_str}) -> {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
                elif "Already compacted" in c_res.stderr or "Already compacted" in c_res.stdout:
                    print(f"  ✓ Already compacted.")
                    skipped += 1
                else:
                    print(f"  ✗ {c_res.stderr.strip()}; retrying with --max-lines 80 fallback...")
                    fb_res = subprocess.run(
                        ["openclaw", "sessions", "compact", key, "--agent", agent, "--max-lines", "80", "--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(5)  # 5s pause to avoid rotational HDD I/O saturation
            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 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=30.0)
            con.execute("PRAGMA optimize;")
            con.close()
        except Exception as e:
            print(f"  ✗ Error optimizing {db}: {e}")
    print("SQLite 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_tracking_sessions()
    compact_tier_a_sessions()
    compact_active_sessions()
    checkpoint_wal_journals()
    vacuum_and_optimize()

if __name__ == "__main__":
    main()