#!/usr/bin/env python3
import glob, sqlite3, json, os, sys, time

t0 = time.time()
recovered = 0
interrupted_turns = []

def extract_prompt(msg):
    if not msg: return ""
    content = msg.get("content")
    if isinstance(content, str): return content
    if isinstance(content, list):
        parts = []
        for c in content:
            if isinstance(c, dict): parts.append(c.get("text") or c.get("content") or "")
            elif isinstance(c, str): parts.append(c)
        return "\n".join(filter(None, parts))
    return str(content or "")

for db in sorted(glob.glob("/home/node/.openclaw/agents/*/agent/openclaw-agent.sqlite")):
    try:
        conn = sqlite3.connect(db, timeout=10.0)
        cursor = conn.cursor()
        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,))
        cursor.execute("SELECT session_key, entry_json, current_session_id, status FROM session_nodes WHERE status IN (\"running\", \"failed\", \"error\")")
        rows = cursor.fetchall()
        agent_id = os.path.basename(os.path.dirname(os.path.dirname(db)))
        for sk, ej, csid, st in rows:
            try:
                p = json.loads(ej) if ej else {}
                sid = csid or p.get("sessionId")
                if sid:
                    ev = cursor.execute("SELECT seq, event_json FROM transcript_events WHERE session_id = ? ORDER BY seq DESC LIMIT 1", (sid,)).fetchone()
                    if ev:
                        ev_seq, ev_raw = ev
                        try:
                            ev_data = json.loads(ev_raw)
                            msg = ev_data.get("message") or ev_data
                            if st == "running" and msg.get("role") == "user":
                                prompt_text = extract_prompt(msg)
                                ctx = p.get("restartRecoveryDeliveryContext") or p.get("delivery", {}).get("context") or p.get("delivery", {}).get("route", {}).get("context") or {}
                                interrupted_turns.append({
                                    "agent": agent_id,
                                    "session_key": sk,
                                    "session_id": sid,
                                    "context": ctx,
                                    "prompt": prompt_text
                                })
                                cursor.execute("PRAGMA foreign_keys = ON")
                                cursor.execute("DELETE FROM transcript_events WHERE session_id = ? AND seq = ?", (sid, ev_seq))
                            elif msg.get("role") == "assistant" and msg.get("stopReason") == "error":
                                cursor.execute("PRAGMA foreign_keys = ON")
                                cursor.execute("DELETE FROM transcript_events WHERE session_id = ? AND seq = ?", (sid, ev_seq))
                        except Exception as ex:
                            print(f"[recover_sessions] Error parsing event: {ex}", file=sys.stderr)
                p["status"] = "done"
                p["activeWriterRunId"] = None
                p["activeClaimRunId"] = None
                p["restartRecoveryDeliveryRunId"] = None
                cursor.execute("UPDATE session_nodes SET status = \"done\", entry_json = ? WHERE session_key = ?", (json.dumps(p), sk))
                cursor.execute("UPDATE session_nodes SET entry_valid = 1 WHERE session_key = ?", (sk,))
                recovered += 1
            except Exception: pass
        cursor.execute("UPDATE session_nodes SET entry_valid = 1 WHERE entry_valid != 1")
        cursor.execute("UPDATE session_windows SET status = \"done\" WHERE status IN (\"running\", \"failed\", \"error\")")
        conn.commit()
        conn.close()
    except Exception as e:
        print(f"[recover_sessions] Warning for {db}: {e}", file=sys.stderr)

state_db = "/home/node/.openclaw/state/openclaw.sqlite"
if os.path.exists(state_db):
    try:
        sconn = sqlite3.connect(state_db, timeout=10.0)
        scur = sconn.cursor()
        scur.execute("UPDATE channel_ingress_events SET attempts = 0, last_attempt_at = NULL, last_error = NULL WHERE status = 'pending' AND attempts > 0")
        sconn.commit()
        sconn.close()
    except Exception as e:
        print(f"[recover_sessions] Warning resetting ingress backoff: {e}", file=sys.stderr)

if interrupted_turns:
    try:
        with open("/tmp/openclaw-interrupted-turns.json", "w") as f:
            json.dump(interrupted_turns, f, indent=2)
        print(f"[recover_sessions] Detected {len(interrupted_turns)} interrupted user turn(s) saved to /tmp/openclaw-interrupted-turns.json")
    except Exception as e:
        print(f"[recover_sessions] Error writing interrupted turns: {e}", file=sys.stderr)

print(f"[recover_sessions] Completed in {time.time() - t0:.2f}s ({recovered} wedged sessions recovered)")