#!/usr/bin/env python3
"""crew-shutdown [PAIR|all] — stop crews safely and say what is outstanding.

It does not just kill things. In order:

  1. Asks every running Coordinator whether anything is waiting on the owner, and
     gives it a moment to answer.
  2. Reads each Coordinator's last `the owner:` line and repeats any real question.
  3. Has every agent write its handoff, so nothing it knows closes with it.
  4. Turns the relay off so nothing new is handed out.
  5. Closes the panes, and clears the files the harness leaves behind.
  6. Reports what the crews are storing, so nothing grows unnoticed.

Nothing is deleted. Everything the crews wrote stays on disk.

    crew-shutdown            every crew that is running
    crew-shutdown <name>     just that one
    crew-shutdown --dry-run  say what it would do, change nothing
"""
import json
import os
import re
import signal
import subprocess
import sys
import time
from pathlib import Path

HOME = Path.home()
CREW = HOME / ".local/share/crew"
PANES = CREW / "panes"
PAIRS = CREW / "pairs.json"
BIN = HOME / "bin"
ROLES = ("coordinator", "architect", "builder")

FRESH = 30          # seconds; a pane file older than this means the crew is gone
ANSWER_WAIT = float(os.environ.get("CREW_SHUTDOWN_WAIT", "45"))  # seconds for a Coordinator to answer
# How long all the agents together get to write their handoffs before the panes go.
WRITE_MINUTES = float(os.environ.get("CREW_SHUTDOWN_HANDOFF_MINUTES", "5"))
DRY = "--dry-run" in sys.argv


def registry():
    try:
        return json.loads(PAIRS.read_text() or "{}")
    except (OSError, ValueError):
        return {}


def panes_alive(pair):
    """Process ids of that crew's panes, whether it runs in the harness or in
    Konsole. Matched on the whole command line so nothing else can match."""
    out = {}
    for role in ROLES:
        cmd = (registry().get(pair, {}).get("ids") or {}).get(role, "")
        if not cmd:
            continue
        pids = subprocess.run(["pgrep", "-f", f"^bash {cmd}$"],
                              capture_output=True, text=True).stdout.split()
        if pids:
            out[role] = [int(x) for x in pids]
    return out


def running(pair):
    """Alive if its panes are, or if the harness is still writing their screens.

    Checking only the screen files missed a crew running in Konsole, which has
    no screen files at all.
    """
    if panes_alive(pair):
        return True
    try:
        f = PANES / f"{pair}-coordinator.screen"
        return time.time() - f.stat().st_mtime < FRESH
    except OSError:
        return False


def screen(pair, role):
    try:
        return (PANES / f"{pair}-{role}.screen").read_text(errors="replace")
    except OSError:
        return ""


def question(pair):
    """The Coordinator's last the owner: line, unless it is the word nothing."""
    for line in reversed(screen(pair, "coordinator").splitlines()):
        m = re.match(r"\s*the owner:\s*(.+)", line)
        if m:
            asked = m.group(1).strip()
            return "" if asked.lower().startswith("nothing") else asked
    return ""


def size(path):
    total = 0
    for root, _dirs, files in os.walk(path):
        for f in files:
            try:
                total += (Path(root) / f).stat().st_size
            except OSError:
                pass
    return total


def gb(n):
    return f"{n / 1_073_741_824:.1f} GB"


def storage_report():
    """What the crews are storing. Nothing here is deleted, only measured."""
    print("\nWhat the crews are storing")
    rows = [("Codex sessions", HOME / ".codex/sessions", 3.0),
            ("Claude transcripts", HOME / ".claude/projects", 2.0),
            ("Crew chat folders", Path(os.environ.get("CREW_DOCS", HOME / "agent-docs")) / "Chats", 1.0)]
    for label, path, limit_gb in rows:
        if not path.exists():
            continue
        n = size(path)
        note = ""
        if n / 1_073_741_824 > limit_gb:
            note = "  LARGE — see the biggest files below"
        print(f"  {label:<20}{gb(n)}{note}")
        if note:
            big = []
            for root, _d, files in os.walk(path):
                for f in files:
                    p = Path(root) / f
                    try:
                        big.append((p.stat().st_size, p))
                    except OSError:
                        pass
            for n2, p in sorted(big, reverse=True)[:3]:
                print(f"      {n2 / 1_048_576:.0f} MB  {p.name[:60]}")


def handoffs(pair):
    """Ask every running agent to write its handoff before its pane closes.

    A pane that closes without one takes everything it knew with it: the next
    opening finds a log that is not from today, starts clean by the rule in
    crew-start-mode, and has nothing to read. Nobody can ask a session that has
    already gone, so the ask belongs here, while the agents are alive.

    Gives back the roles that wrote none. It never stops the shutdown — he asked
    for the crews to close, and a silent agent is reported, not obeyed.
    """
    if not running(pair):
        return []
    print(f"  asking each agent to write its handoff (up to {WRITE_MINUTES:g} min)")
    if DRY:
        return []
    r = subprocess.run(
        [str(BIN / "crew-handoff"), "--write-only", f"--pair={pair}"],
        capture_output=True, text=True,
        env={**os.environ, "CREW_HANDOFF_WRITE_MINUTES": str(WRITE_MINUTES)})
    wrote = re.findall(r"(\w+): handoff written \((.+)\)$", r.stdout, re.M)
    missing = re.findall(r"(\w+): no handoff written", r.stdout, re.M)
    for role, note in wrote:
        print(f"    {role}: {note}")
    for role in missing:
        print(f"    {role}: NO handoff — its next session starts with nothing")
    if not wrote and not missing:
        print("    no agent answered; nothing was written")
    return missing


def stop(pair):
    name = registry().get(pair, {}).get("name", pair)
    print(f"\n{name}")
    asked = ""
    if running(pair):
        print("  asking the Coordinator if anything is waiting on you")
        if not DRY:
            subprocess.run(
                [str(BIN / "crew-say"), pair, "coordinator",
                 "Shutting down. Is anything waiting on the owner? Answer in his four"
                 " blocks and end with the the owner: line. Do not start new work."],
                capture_output=True, text=True)
            time.sleep(ANSWER_WAIT)
        asked = question(pair)
        print("  waiting on you: " + (asked if asked else "nothing"))
    else:
        print("  not running")

    handoffs(pair)

    if not DRY:
        subprocess.run([str(BIN / "pqpw"), "auto", "off", pair],
                       capture_output=True, text=True)
    print("  relay off")

    # A crew in the harness is closed by the harness itself; stopping its panes
    # from outside leaves the page alive, writing its files straight back.
    if not DRY:
        try:
            PANES.mkdir(parents=True, exist_ok=True)
            (PANES / f"{pair}.close").write_text("close\n")
        except OSError:
            pass
        for _ in range(20):
            if not (PANES / f"{pair}.close").exists():
                time.sleep(2)
                break
            time.sleep(0.5)
        try:
            (PANES / f"{pair}.close").unlink()
        except OSError:
            pass

    alive = panes_alive(pair)
    closed, orphans = 0, 0
    for role, pids in alive.items():
        # More than one process for a role means earlier copies were orphaned —
        # a pane whose window went without taking its shell with it.
        if len(pids) > 1:
            orphans += len(pids) - 1
        for pid in pids:
            closed += 1
            if not DRY:
                try:
                    os.kill(pid, signal.SIGTERM)
                except OSError:
                    pass
    print(f"  panes closed: {closed}"
          + (f" ({orphans} of them orphaned from earlier openings)" if orphans else ""))

    if not DRY:
        for role in ROLES:
            for suffix in (".screen", ".in"):
                try:
                    (PANES / f"{pair}-{role}{suffix}").unlink()
                except OSError:
                    pass
    print("  leftover files cleared")
    return asked


def main():
    args = [a for a in sys.argv[1:] if not a.startswith("-")]
    crews = registry()
    targets = [args[0]] if args and args[0] != "all" else list(crews)
    unknown = [t for t in targets if t not in crews]
    if unknown:
        sys.exit("no crew called " + ", ".join(unknown))

    print("Shutting the crews down safely" + (" (dry run)" if DRY else ""))
    questions = {}
    for pair in sorted(targets):
        asked = stop(pair)
        if asked:
            questions[crews[pair].get("name", pair)] = asked

    storage_report()

    print("\nWhat needs you")
    if questions:
        for name, asked in questions.items():
            print(f"  {name}: {asked}")
    else:
        print("  nothing")
    print("\nNothing was deleted. Every handoff, report and round log is on disk.")
    return 0


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