#!/usr/bin/env python3
"""omni-cswap rescue watchdog: auto-recover Omnigent claude sessions that hit a usage limit.

Loop: scan Omnigent's tmux panes → detect a rate/usage-limit error (or login screen) in the
pane text → exclude that account and re-pick a healthy one (writes the sticky assignment the
shim will honor) → respawn the pane in place (same command → shim → new account, `--resume`
appended so the conversation cold-resumes) → re-send the last user message.

Deliberately tmux-based, not Omnigent-API-based: the API bearer token expires every ~8h,
tmux + transcripts never do. Runs fine as a LaunchAgent / systemd --user service.

Env knobs: OMNI_RESCUE_INTERVAL (30s) · OMNI_RESCUE_COOLDOWN (600s) · OMNI_RESCUE_GLOBS
(colon-separated extra socket globs, used by tests) · OMNI_RESCUE_ONCE=1 (single tick)."""
import glob, json, os, re, shlex, subprocess, sys, time
from pathlib import Path

SELF = Path(__file__).resolve().parent
LOG = SELF / "rescue.log"
STATE = SELF / "rescue-state.json"
INTERVAL = float(os.environ.get("OMNI_RESCUE_INTERVAL", "30"))
COOLDOWN = float(os.environ.get("OMNI_RESCUE_COOLDOWN", "600"))
CLAUDE_PROJECTS = Path.home() / ".claude" / "projects"

LIMIT_RE = re.compile(r"usage limit|rate limit|reached your|weekly limit|hit your|/usage-credits|switch models with /model|out of extended usage|limit (?:will )?reset", re.I)
LOGIN_RE = re.compile(r"select login method", re.I)
READY_RE = re.compile(r"\? for shortcuts|shift\+tab to cycle|bypass permissions", re.I)
# One-time confirmation gates that stall an autonomous session (they persist per config dir once
# accepted): the Bypass-Permissions-mode warning, and the folder-trust dialog. Auto-accept them.
GATE_RE = re.compile(r"Bypass Permissions mode|trust this folder|Quick safety check", re.I)
GATE_ACCEPT_RE = re.compile(r"(?:❯\s*)?(\d+)\.\s*(?:Yes, I accept|Yes, I trust this folder)")

def log(m):
    line = f"{time.strftime('%FT%TZ', time.gmtime())} {m}"
    print(line, flush=True)
    with open(LOG, "a") as f: f.write(line + "\n")

def tmux(sock, *args, timeout=10):
    return subprocess.run(["tmux", "-S", sock, *args], capture_output=True, text=True, timeout=timeout)

def sockets():
    globs = ["/tmp/omnigent-terminal-*/tmux.sock"]
    tmp = os.environ.get("TMPDIR", "")
    if tmp: globs.append(os.path.join(tmp, "omnigent-terminal-*", "tmux.sock"))
    globs += ["/var/folders/*/*/T/omnigent-terminal-*/tmux.sock"]
    globs += [g for g in os.environ.get("OMNI_RESCUE_GLOBS", "").split(":") if g]
    out = set()
    for g in globs: out.update(glob.glob(g))
    return sorted(out)

def panes(sock):
    r = tmux(sock, "list-panes", "-aF", "#{pane_id}\t#{pane_pid}\t#{pane_current_path}\t#{pane_start_command}")
    if r.returncode != 0: return []
    rows = []
    for line in r.stdout.splitlines():
        parts = line.split("\t", 3)
        if len(parts) == 4: rows.append(parts)
    return rows

def pane_text(sock, pane):
    r = tmux(sock, "capture-pane", "-p", "-t", pane, "-S", "-120")
    return r.stdout if r.returncode == 0 else ""

def env_of(pid):
    try:
        if sys.platform == "linux":
            data = open(f"/proc/{pid}/environ", "rb").read().split(b"\0")
            return dict(kv.decode(errors="replace").split("=", 1) for kv in data if b"=" in kv)
        out = subprocess.run(["ps", "eww", str(pid)], capture_output=True, text=True, timeout=5).stdout
        return dict(tok.split("=", 1) for tok in out.split() if "=" in tok and tok.split("=", 1)[0].isupper())
    except Exception:
        return {}

def account_of(pid):
    cfg = env_of(pid).get("CLAUDE_CONFIG_DIR", "")
    m = re.search(r"/sessions/(\d+)-", cfg)
    if m: return m.group(1)
    try:
        out = subprocess.run(["cswap", "status", "--json"], capture_output=True, text=True, timeout=20).stdout
        return str(json.loads(out).get("activeAccountNumber", "")) or None
    except Exception:
        return None

def transcript_for(cwd):
    d = CLAUDE_PROJECTS / re.sub(r"[^A-Za-z0-9]", "-", os.path.realpath(cwd))
    if not d.is_dir(): return None
    files = sorted(d.glob("*.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True)
    return files[0] if files else None

def last_user_message(tr):
    try: lines = tr.read_text(errors="replace").splitlines()
    except Exception: return None
    for line in reversed(lines):
        try: obj = json.loads(line)
        except Exception: continue
        if obj.get("type") != "user" or obj.get("isMeta"): continue
        msg = obj.get("message") or {}
        c = msg.get("content")
        text = ""
        if isinstance(c, str): text = c
        elif isinstance(c, list):
            text = "\n".join(b.get("text", "") for b in c if isinstance(b, dict) and b.get("type") == "text")
        text = text.strip()
        if text and not text.startswith("<") and "tool_result" not in str(msg)[:200]:
            return text[:4000]
    return None

def model_of(start_cmd):
    toks = shlex.split(start_cmd) if start_cmd else []
    for i, t in enumerate(toks):
        if t == "--model" and i + 1 < len(toks): return toks[i + 1]
        if t.startswith("--model="): return t.split("=", 1)[1]
    return ""

def repick(cwd, exclude, model=""):
    env = dict(os.environ, OMNI_CSWAP_EXCLUDE=exclude, OMNI_CSWAP_MODEL=model)
    env.pop("OMNIGENT_SESSION_ID", None)
    # also clear the old sticky so the pick is fresh
    try:
        st = json.loads((SELF / "assignments.json").read_text())
        st.pop(f"cwd:{os.path.realpath(cwd)}", None)
        (SELF / "assignments.json").write_text(json.dumps(st, indent=1))
    except Exception: pass
    r = subprocess.run(["python3", str(SELF / "pick_account.py")], capture_output=True, text=True, cwd=cwd, env=env, timeout=45)
    return r.stdout.strip() or None

def rescue(sock, pane, pid, cwd, start_cmd, reason):
    acct = account_of(pid)
    tr = transcript_for(cwd)
    resume_id = tr.stem if tr else None
    msg = last_user_message(tr) if tr else None
    new = repick(cwd, acct or "", model_of(start_cmd))
    if not new:
        log(f"RESCUE ABORT {pane}@{sock}: no healthy alternative account (failed acct {acct})")
        return False
    # Rebuild the claude invocation robustly: strip any leading `VAR=val` env prefixes that
    # tmux's pane_start_command may or may not preserve, then run it under an explicit PATH that
    # includes the shim + stable claude dirs (Mac ~/.local/bin, Linux fnm) — the shim re-derives
    # everything from PATH, and the sticky we just wrote (excluding the capped account) routes it.
    toks = shlex.split(start_cmd) if start_cmd else []
    while toks and re.match(r"^[A-Za-z_][A-Za-z0-9_]*=", toks[0]):
        toks.pop(0)
    if not toks:
        toks = ["claude"]
    if resume_id and "--resume" not in toks:
        toks += ["--resume", resume_id]
    home = str(Path.home())
    good_path = ":".join([str(SELF), f"{home}/.local/bin", f"{home}/.fnm/aliases/default/bin",
                          "/opt/homebrew/bin", "/opt/homebrew/sbin", "/usr/bin", "/bin"])
    cmd = "env " + shlex.quote(f"PATH={good_path}") + " " + " ".join(shlex.quote(t) for t in toks)
    log(f"RESCUE {pane}@{os.path.basename(os.path.dirname(sock))}: reason={reason} acct {acct} -> {new}; respawn (resume={resume_id}, resend={'yes' if msg else 'no'})")
    r = tmux(sock, "respawn-pane", "-k", "-t", pane, cmd)
    if r.returncode != 0:
        log(f"  respawn failed: {r.stderr.strip()[:200]}"); return False
    if msg:
        deadline = time.time() + 90
        while time.time() < deadline:
            time.sleep(5)
            if READY_RE.search(pane_text(sock, pane)): break
        time.sleep(3)
        tmux(sock, "send-keys", "-t", pane, "-l", msg)
        time.sleep(1)
        tmux(sock, "send-keys", "-t", pane, "Enter")
        log(f"  resent last user message ({len(msg)} chars)")
    return True

def tick(state):
    now = time.time()
    for sock in sockets():
        for pane, pid, cwd, start_cmd in panes(sock):
            key = f"{sock}:{pane}"
            text = pane_text(sock, pane)
            tail = "\n".join(text.splitlines()[-40:])
            # First: auto-accept one-time gates (bypass-mode warning / folder trust). No cooldown —
            # these must clear immediately, and they self-limit (persist after one accept).
            # Search the FULL pane text, not just the tail: the gate warning renders at the TOP.
            if GATE_RE.search(text):
                m = GATE_ACCEPT_RE.search(text)
                if m:
                    tmux(sock, "send-keys", "-t", pane, m.group(1)); time.sleep(0.4)
                    tmux(sock, "send-keys", "-t", pane, "Enter")
                    log(f"auto-accepted gate ({m.group(1)}) on {pane}@{os.path.basename(os.path.dirname(sock))}")
                continue
            if now - state.get(key, 0) < COOLDOWN: continue
            reason = "limit" if LIMIT_RE.search(tail) else ("login" if LOGIN_RE.search(tail) else None)
            if not reason: continue
            state[key] = now
            STATE.write_text(json.dumps(state))
            try: rescue(sock, pane, pid, cwd, start_cmd, reason)
            except Exception as e: log(f"rescue error {key}: {e}")

def main():
    try: state = json.loads(STATE.read_text())
    except Exception: state = {}
    log(f"rescue watchdog start (interval={INTERVAL}s cooldown={COOLDOWN}s)")
    while True:
        try: tick(state)
        except Exception as e: log(f"tick error: {e}")
        if os.environ.get("OMNI_RESCUE_ONCE"): return
        time.sleep(INTERVAL)

if __name__ == "__main__":
    main()