#!/usr/bin/env python3
"""apply — apply amux-nuilab's fixes to your local copy of amux-server.py.

These are small, surgical fixes to upstream amux (github.com/mixpeek/amux) that
have not landed upstream. Nothing here redistributes upstream code: each patch is
an anchored edit applied to the copy already installed on your machine.

  running-flag   `stop` is a GRACEFUL stop: it exits the agent but deliberately
                 leaves the tmux shell alive. The session list computes
                 `running = tmux_name(name) in tmux_info`, i.e. only "does a tmux
                 session exist" — so a stopped session reads as running forever in
                 the dashboard, in /api/sessions, and in `amux-remote <box> ls`.
                 It also disagrees with the server's own is_running(), which does
                 check whether the pane is sitting at a bare shell prompt.
                 FIX: apply that same shell-prompt test in the list builder. No
                 extra tmux calls — the pane capture is already in hand.

  auto-resume    When start_session cannot resolve a conversation to resume (stale
                 uuid, ambiguous name, missing meta) it falls back to a BLANK
                 `--name` start. The watchdog restarts sessions automatically, so
                 that silently wipes the session's whole context — and each wipe
                 writes another same-titled transcript, which makes the next
                 resume MORE ambiguous. A self-worsening loop.
                 FIX: before starting blank, ask upstream's own
                 _find_latest_session_id() for the newest real conversation in
                 that directory and resume it. Upstream already ships that helper;
                 it just is not consulted on the fallback paths.

Safety properties, because this edits a 3 MB file you did not write:
  * idempotent      — a marker comment means re-running is a no-op
  * anchored        — refuses if the expected code is not found (upstream drifted)
  * fail-closed     — validates the result with py_compile BEFORE replacing anything
  * reversible      — backs up first; --revert restores the newest backup

Usage:
    patches/apply              apply every patch that is missing
    patches/apply --check      report status, change nothing (exit 1 if unpatched)
    patches/apply --revert     restore the most recent backup
    patches/apply --server P   operate on a specific amux-server.py
"""
import argparse
import datetime
import os
import py_compile
import re
import shutil
import sys
import tempfile

MARKER = "amux-nuilab patch"

SEARCH_PATHS = [
    "~/.local/bin/amux-server.py",
    "/usr/local/bin/amux-server.py",
    "/opt/amux/amux-server.py",
    "~/.local/share/amux/amux-server.py",
]


def find_server(explicit=None):
    if explicit:
        p = os.path.expanduser(explicit)
        if not os.path.isfile(p):
            sys.exit(f"no such file: {p}")
        return p
    for c in SEARCH_PATHS:
        p = os.path.expanduser(c)
        if os.path.isfile(p):
            return p
    # Fall back to whatever is next to the amux on PATH.
    amux = shutil.which("amux")
    if amux:
        cand = os.path.join(os.path.dirname(os.path.realpath(amux)), "amux-server.py")
        if os.path.isfile(cand):
            return cand
    sys.exit(
        "could not find amux-server.py.\n"
        "  Looked in: " + ", ".join(SEARCH_PATHS) + "\n"
        "  Pass it explicitly:  patches/apply --server /path/to/amux-server.py"
    )


# ── patch 1: truthful running flag ───────────────────────────────────────────
def patch_running_flag(text):
    """Downgrade running->False when the pane is at a bare shell prompt.

    Anchored on the list builder's preview/status block. The status helper has
    been renamed upstream before (_detect_claude_status -> _detect_session_status),
    so match it loosely rather than pinning the exact call.
    """
    if "amux-nuilab patch: graceful stop leaves the tmux shell alive" in text:
        return text, "already applied"

    pat = re.compile(
        r'( *)preview = strip_ansi\(lines\[-1\]\[:120\]\) if lines else ""\n'
        r'( *)if running:\n'
        r'( *)status = _detect_\w+\(',
    )
    m = pat.search(text)
    if not m:
        return text, "ANCHOR NOT FOUND — upstream changed; not applied"

    indent = m.group(1)
    insert = (
        f'{indent}# {MARKER}: graceful stop leaves the tmux shell alive, so "a tmux\n'
        f'{indent}# session exists" is NOT the same as "the agent is running". Apply the\n'
        f'{indent}# same bare-shell-prompt test is_running() uses, or a gracefully stopped\n'
        f'{indent}# session reports running forever in the list and in /api/sessions.\n'
        f'{indent}if running and _at_shell_prompt(strip_ansi(raw)):\n'
        f'{indent}    running = False\n'
    )
    start = m.start(2)
    return text[:start] + insert + text[start:], "applied"


# ── patch 2: never blank-start over a real conversation ──────────────────────
def patch_auto_resume(text):
    """Resume the newest real conversation instead of starting blank.

    Every fallback in start_session assigns `session_flag = f'--name ...'`. Rather
    than rewrite each branch, add one safety net just after the whole resolution
    block: if we ended up on a blank --name start but a real conversation exists
    for this directory, resume that instead.
    """
    if "amux-nuilab patch: never blank-start over a real conversation" in text:
        return text, "already applied"

    if "def _find_latest_session_id" not in text:
        return text, "upstream helper _find_latest_session_id missing — not applied"

    # Anchor: the LAST fallback in the chain, the first-ever-start branch.
    anchor = (
        "            else:\n"
        "                # First-ever start\n"
        "                session_flag = f'--name {shlex.quote(name)}'\n"
    )
    if anchor not in text:
        return text, "ANCHOR NOT FOUND — upstream changed; not applied"
    if text.count(anchor) != 1:
        return text, "anchor is ambiguous — not applied"

    addition = anchor + (
        "\n"
        "            # " + MARKER + ": never blank-start over a real conversation.\n"
        "            # A blank --name start WIPES the session's context, and the watchdog\n"
        "            # restarts sessions on its own — so an unresolved resume target used to\n"
        "            # silently destroy work. Each wipe also wrote another same-titled\n"
        "            # transcript, making the next resume more ambiguous: a self-worsening\n"
        "            # loop. If a real conversation exists for this directory, resume it.\n"
        "            if session_flag.startswith('--name'):\n"
        "                _latest = _find_latest_session_id(work_dir)\n"
        "                if _latest:\n"
        "                    session_flag = f'--resume {_latest}'\n"
        "                    print(f\"[start] {name}: resume newest conversation \"\n"
        "                          f\"(uuid={_latest}) instead of a blank start\")\n"
    )
    return text.replace(anchor, addition, 1), "applied"


PATCHES = [
    ("running-flag", patch_running_flag),
    ("auto-resume", patch_auto_resume),
]


def is_applied(text, name):
    return {
        "running-flag": "graceful stop leaves the tmux shell alive" in text,
        "auto-resume": "never blank-start over a real conversation" in text,
    }[name]


def main():
    ap = argparse.ArgumentParser(add_help=False)
    ap.add_argument("--check", action="store_true")
    ap.add_argument("--revert", action="store_true")
    ap.add_argument("--server")
    ap.add_argument("-h", "--help", action="store_true")
    args = ap.parse_args()

    if args.help:
        print(__doc__)
        return 0

    server = find_server(args.server)
    text = open(server, encoding="utf-8").read()

    if args.check:
        print(f"amux-server.py: {server}")
        missing = 0
        for name, _ in PATCHES:
            if is_applied(text, name):
                print(f"  [applied] {name}")
            else:
                print(f"  [MISSING] {name}")
                missing += 1
        return 1 if missing else 0

    if args.revert:
        backups = sorted(
            f for f in os.listdir(os.path.dirname(server) or ".")
            if f.startswith(os.path.basename(server) + ".amux-nuilab-bak-")
        )
        if not backups:
            sys.exit("no backup to revert to")
        newest = os.path.join(os.path.dirname(server), backups[-1])
        shutil.copy2(newest, server)
        print(f"reverted {server} from {newest}")
        return 0

    if not os.access(server, os.W_OK):
        sys.exit(f"{server} is not writable by you — re-run with the right permissions")

    new = text
    results = []
    for name, fn in PATCHES:
        new, status = fn(new)
        results.append((name, status))

    if new == text:
        for name, status in results:
            print(f"  {name}: {status}")
        print("nothing to do.")
        return 0

    # Validate BEFORE touching the installed file: a 3 MB server that no longer
    # compiles is a much worse outcome than an unapplied patch.
    tmp = tempfile.NamedTemporaryFile("w", suffix=".py", delete=False, encoding="utf-8")
    tmp.write(new)
    tmp.close()
    compiled = tmp.name + ".pyc"
    try:
        # Give py_compile an explicit temporary destination. Some macOS Python
        # builds redirect caches into ~/Library/Caches, which can be unwritable
        # in a sandbox even though the target server and temp directory are fine.
        py_compile.compile(tmp.name, cfile=compiled, doraise=True)
    except py_compile.PyCompileError as e:
        os.unlink(tmp.name)
        if os.path.exists(compiled):
            os.unlink(compiled)
        sys.exit(f"patched file does not compile — refusing to install it.\n{e}")
    if os.path.exists(compiled):
        os.unlink(compiled)

    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    backup = f"{server}.amux-nuilab-bak-{stamp}"
    shutil.copy2(server, backup)
    shutil.copy2(tmp.name, server)
    os.unlink(tmp.name)

    for name, status in results:
        print(f"  {name}: {status}")
    print(f"\nbacked up to {backup}")
    print("restart the server for this to take effect:")
    print("  setup/serve-install        (or just restart `amux serve`)")
    return 0


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