"""Issue #45 (B): close the restored empty shells of the session being resumed.

With restore-on-restart off (lib_py/herdr_config.py) a restarted MAM server
still brings every saved pane back - as an empty shell in its old cwd, or in
$HOME when that directory is gone. A pane keeps the `agent_session` its agent
reported before the server died. resume_session.sh calls this once, right
before it spawns, with the id it is about to resume; a pane is closed only when
all of these hold:

  1. the target is an isolated MAM server (HERDR_SESSION_NAME set, not default);
  2. herdr sees no live agent in it;
  3. its agent_session value is that id - so a user's pane, another session's
     pane and a freshly split pane (no agent_session yet) are never candidates;
  4. its foreground process is the shell itself (nothing running in it).

The registry is not consulted: a session whose server died stays `running`
there, so "running" cannot tell a live pane from a restored one.

CLI: `python3 -m lib_py.herdr_restore close-restored-shells --session-id ID
[--herdr PATH]`. One stderr line per closed pane; exit 1 when herdr cannot be
queried (the caller only warns).
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from typing import List, Optional

# A pane exports these; they would point the query at the caller's own server
# and workspace instead of the --session named here.
_SCOPE_ENV = ("HERDR_WORKSPACE_ID", "HERDR_PANE_ID", "HERDR_TAB_ID", "HERDR_SOCKET_PATH")


def _herdr(herdr: str, session: str, *args: str) -> dict:
    env = {k: v for k, v in os.environ.items() if k not in _SCOPE_ENV}
    out = subprocess.run([herdr, "--session", session, *args], env=env,
                         capture_output=True, text=True, timeout=30, check=True).stdout
    return json.loads(out).get("result") or {}


def _shell_in_foreground(herdr: str, session: str, pane_id: str) -> bool:
    info = _herdr(herdr, session, "pane", "process-info", "--pane", pane_id).get("process_info") or {}
    shell = info.get("shell_pid")
    return shell is not None and info.get("foreground_process_group_id") == shell


def close_restored_shells(session_id: str, herdr: str = "herdr",
                          session: Optional[str] = None) -> List[str]:
    """Close the matching panes; return their ids."""
    session = os.environ.get("HERDR_SESSION_NAME", "") if session is None else session
    if not session_id or session in ("", "default"):
        return []
    closed = []
    for p in _herdr(herdr, session, "pane", "list").get("panes") or []:
        pane_id = p.get("pane_id")
        if not pane_id or p.get("agent"):
            continue
        if (p.get("agent_session") or {}).get("value") != session_id:
            continue
        if not _shell_in_foreground(herdr, session, pane_id):
            continue
        _herdr(herdr, session, "pane", "close", pane_id)
        print(f"closed restored pane {pane_id} (cwd={p.get('cwd', '')}, agent_session={session_id})",
              file=sys.stderr)
        closed.append(pane_id)
    return closed


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description="close a resumed session's restored empty shells (Issue #45)")
    sub = p.add_subparsers(dest="cmd", required=True)
    c = sub.add_parser("close-restored-shells")
    c.add_argument("--session-id", required=True)
    c.add_argument("--herdr", default="", help="the real herdr binary (not the MAM shim)")
    args = p.parse_args(argv)
    try:
        close_restored_shells(args.session_id, args.herdr or "herdr")
    except (OSError, subprocess.SubprocessError, ValueError) as e:
        print(f"close-restored-shells: {e}", file=sys.stderr)
        return 1
    return 0


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