#!/usr/bin/env python3
"""Resume Claude Code background sessions that stalled on a recoverable error.

Claude Code blocks and exits when it hits a usage limit, an API outage, or a
network failure. The process is gone, so nothing scheduled inside the session
can revive it. This runs outside every session and resumes the ones whose
cause has cleared.
"""

import argparse
import fcntl
import json
import os
import re
import shutil
import subprocess
import sys
from datetime import datetime, timedelta
from pathlib import Path
from zoneinfo import ZoneInfo

# Where `claude` installs itself, and what a scheduler's minimal PATH omits.
CLAUDE_FALLBACK_DIRS = ("~/.local/bin", "/usr/local/bin", "/opt/homebrew/bin")

DEFAULT_PROMPT = (
    "Continue where you left off. Your previous run was interrupted by a "
    "transient failure (usage limit, API outage, or network), which has cleared."
)

RECOVERABLE = re.compile(
    r"hit your (session|weekly|opus) limit"
    r"|usage limit"
    r"|rate.?limit"
    r"|api error"
    r"|overloaded"
    r"|connection (error|reset|closed|failure)"
    r"|network error"
    r"|fetch failed"
    r"|socket hang up"
    r"|etimedout|econnreset|econnrefused|enotfound",
    re.IGNORECASE,
)

# A session parked on a human decision must never be auto-resumed. Vetoes even
# when the same detail also names a limit.
NEEDS_HUMAN = re.compile(
    r"input needed|waiting for|permission|approval"
    r"|awaiting (your )?(direction|decision|input|answer|go.?ahead)"
    r"|needs? (your )?(direction|decision|input)",
    re.IGNORECASE,
)

LIMIT_BLOCK = re.compile(r"(session|weekly|opus|usage|rate.?) ?limit", re.IGNORECASE)
WEEKLY_LIMIT = re.compile(r"weekly limit", re.IGNORECASE)
SESSION_WINDOW = timedelta(hours=5)
WEEKLY_WINDOW = timedelta(days=7)

RESETS_AT = re.compile(
    r"resets\s+(?:(?P<dow>mon|tue|wed|thu|fri|sat|sun)[a-z]*\s+)?"
    r"(?P<hour>\d{1,2})(?::(?P<minute>\d{2}))?\s*(?P<meridiem>am|pm)?"
    r"(?:\s*\((?P<tz>[A-Za-z_]+/[A-Za-z_]+)\))?",
    re.IGNORECASE,
)

WEEKDAYS = {"mon": 0, "tue": 1, "wed": 2, "thu": 3, "fri": 4, "sat": 5, "sun": 6}


def log(message):
    print(f"{datetime.now().astimezone().isoformat(timespec='seconds')} {message}", flush=True)


def resolve_claude(name):
    """Absolute path to the Claude Code binary, or None if it cannot be found."""
    found = shutil.which(name)
    if found:
        return found
    for directory in CLAUDE_FALLBACK_DIRS:
        candidate = Path(directory).expanduser() / name
        if candidate.is_file() and os.access(candidate, os.X_OK):
            return str(candidate)
    return None


def single_instance(path):
    """Exclusive lock for this run, or None when another run already holds it.

    Overlapping runs both read the ledger before either writes it, so both
    resume the same session — two agents racing on one git worktree.
    """
    path.parent.mkdir(parents=True, exist_ok=True)
    handle = path.open("w")
    try:
        fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        handle.close()
        return None
    return handle


def active_sessions(claude_bin):
    """Sessions Claude Code still considers active, or None if they cannot be read."""
    try:
        done = subprocess.run(
            [claude_bin, "agents", "--json"],
            capture_output=True, text=True, timeout=120, check=True,
        )
    except (OSError, subprocess.SubprocessError) as err:
        log(f"error: could not list sessions: {err}")
        return None
    try:
        return json.loads(done.stdout)
    except json.JSONDecodeError as err:
        log(f"error: unparseable session list: {err}")
        return None


def pid_alive(pid):
    if not isinstance(pid, int):
        return False
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    return True


def parse_reset(detail, blocked_at):
    """The reset moment named in `detail`, or None. `blocked_at` anchors day-less times."""
    found = RESETS_AT.search(detail)
    if not found:
        return None
    hour = int(found.group("hour"))
    minute = int(found.group("minute") or 0)
    meridiem = (found.group("meridiem") or "").lower()
    if meridiem == "pm" and hour != 12:
        hour += 12
    elif meridiem == "am" and hour == 12:
        hour = 0
    if hour > 23 or minute > 59:
        return None

    zone = None
    if found.group("tz"):
        try:
            zone = ZoneInfo(found.group("tz"))
        except Exception:
            zone = None

    anchor = blocked_at.astimezone(zone) if zone else blocked_at
    target = anchor.replace(hour=hour, minute=minute, second=0, microsecond=0)
    day = (found.group("dow") or "").lower()[:3]
    if day in WEEKDAYS:
        target += timedelta(days=(WEEKDAYS[day] - target.weekday()) % 7)
        if target < anchor:
            target += timedelta(days=7)
    elif target < anchor:
        target += timedelta(days=1)
    return target


def reset_moment(detail, blocked_at):
    """When the block should have lifted, or None if it can be retried at once.

    A stated time is clamped to the limit window: state.json is sometimes
    rewritten after the reset already happened, which would otherwise push the
    parsed time-of-day a full day into the future.
    """
    if not LIMIT_BLOCK.search(detail):
        return None
    window = WEEKLY_WINDOW if WEEKLY_LIMIT.search(detail) else SESSION_WINDOW
    stated = parse_reset(detail, blocked_at)
    bound = blocked_at + window
    return min(stated, bound) if stated else bound


def load_ledger(path):
    try:
        ledger = json.loads(path.read_text())
    except (OSError, json.JSONDecodeError):
        ledger = {}
    ledger.setdefault("attempts", {})
    ledger.setdefault("resumes", [])
    return ledger


def save_ledger(path, ledger):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(ledger, indent=2, sort_keys=True))


def candidates(sessions, jobs_dir, now):
    """Stalled background sessions, oldest block first."""
    found = []
    for session in sessions:
        job_id, session_id = session.get("id"), session.get("sessionId")
        if session.get("kind") != "background" or not job_id or not session_id:
            continue
        if pid_alive(session.get("pid")):
            continue

        state_file = jobs_dir / job_id / "state.json"
        try:
            state = json.loads(state_file.read_text())
        except (OSError, json.JSONDecodeError):
            continue
        if state.get("state") != "blocked":
            continue

        detail = str(state.get("detail") or "")
        if NEEDS_HUMAN.search(detail) or not RECOVERABLE.search(detail):
            continue

        blocked_at = datetime.fromtimestamp(state_file.stat().st_mtime).astimezone()
        reset_at = reset_moment(detail, blocked_at)
        if reset_at and reset_at > now:
            log(f"{job_id}: waiting for reset at {reset_at.isoformat(timespec='minutes')}")
            continue

        found.append({
            "cwd": session.get("cwd"),
            "detail": detail,
            "id": job_id,
            "blocked_at": blocked_at,
            "name": session.get("name"),
            "session_id": session_id,
        })
    return sorted(found, key=lambda job: job["blocked_at"])


def resume(job, claude_bin, prompt, dry_run):
    cwd = job["cwd"]
    if not cwd or not Path(cwd).is_dir():
        log(f"{job['id']}: skipped, working directory is gone ({cwd})")
        return False
    if dry_run:
        log(f"{job['id']}: would resume '{job['name']}' in {cwd} — {job['detail']}")
        return False
    command = [claude_bin, "--bg", "--resume", job["session_id"]]
    if job["name"]:
        command += ["--name", job["name"]]
    try:
        done = subprocess.run(
            command + [prompt],
            capture_output=True, cwd=cwd, text=True, timeout=180, check=True,
        )
    except (OSError, subprocess.SubprocessError) as err:
        log(f"{job['id']}: resume failed: {err}")
        return False
    log(f"{job['id']}: resumed '{job['name']}' — {job['detail']}")

    # The husk stays 'blocked' and active forever otherwise, and would be
    # resumed again into a duplicate session working the same conversation.
    try:
        subprocess.run(
            [claude_bin, "stop", job["id"]],
            capture_output=True, text=True, timeout=60, check=True,
        )
    except (OSError, subprocess.SubprocessError) as err:
        log(f"{job['id']}: could not retire stale entry: {err}")
    return True


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--claude-bin", default=os.environ.get("CLAUDE_BIN", "claude"))
    parser.add_argument(
        "--claude-home",
        type=Path,
        default=Path(os.environ.get("CLAUDE_HOME", Path.home() / ".claude")),
    )
    parser.add_argument("--dry-run", action="store_true", help="report candidates, resume nothing")
    parser.add_argument(
        "--max-attempts", type=int,
        default=int(os.environ.get("WATCHDOG_MAX_ATTEMPTS", "1")),
        help="resumes of any one session, ever (default: 1; its successor is "
             "eligible on its own if it stalls again)",
    )
    parser.add_argument(
        "--max-per-day", type=int,
        default=int(os.environ.get("WATCHDOG_MAX_PER_DAY", "50")),
        help="runaway backstop across all sessions per 24h (default: 50)",
    )
    parser.add_argument(
        "--max-per-run", type=int,
        default=int(os.environ.get("WATCHDOG_MAX_PER_RUN", "0")),
        help="resumes per invocation, 0 for every eligible session (default: 0)",
    )
    parser.add_argument("--prompt", default=os.environ.get("WATCHDOG_PROMPT", DEFAULT_PROMPT))
    args = parser.parse_args()

    claude_bin = resolve_claude(args.claude_bin)
    if not claude_bin:
        log(f"error: '{args.claude_bin}' is not in PATH or "
            f"{', '.join(CLAUDE_FALLBACK_DIRS)} — resuming nothing")
        return 1

    state_dir = args.claude_home / "session-watchdog"
    lock = None
    if not args.dry_run:
        lock = single_instance(state_dir / ".lock")  # released when main returns
        if lock is None:
            log("another run is still going, skipping")
            return 0

    now = datetime.now().astimezone()
    ledger_path = state_dir / "ledger.json"
    ledger = load_ledger(ledger_path)

    recent = [at for at in ledger["resumes"] if now.timestamp() - at < 86400]
    budget = args.max_per_day - len(recent)
    if args.max_per_run > 0:
        budget = min(budget, args.max_per_run)
    if budget <= 0:
        log(f"circuit breaker: {len(recent)} resumes in the last 24h (max {args.max_per_day})")
        return 0

    sessions = active_sessions(claude_bin)
    if sessions is None:
        return 1

    resumed = 0
    for job in candidates(sessions, args.claude_home / "jobs", now):
        if resumed >= budget:
            break
        tried = ledger["attempts"].get(job["session_id"], 0)
        if tried >= args.max_attempts:
            continue
        if not resume(job, claude_bin, args.prompt, args.dry_run):
            continue
        ledger["attempts"][job["session_id"]] = tried + 1
        ledger["resumes"] = recent + [now.timestamp()]
        recent = ledger["resumes"]
        save_ledger(ledger_path, ledger)
        resumed += 1

    if not resumed:
        log("nothing to resume")
    return 0


if __name__ == "__main__":
    sys.exit(main())