"""Decisions behind the multi-agent-mux-reset skill (Issue #78).

reset.sh keeps the shell sequence (arg parsing, terminal check, update
admission, loop lock, sequential stop/create invocations, exit codes) and
delegates target selection, prechecks, atomic spec-file persistence, 5-state
progress classification, and result-table formatting to this module.
Standard library only, Python 3.9 compatible.
"""
from __future__ import annotations

import argparse
import datetime
import json
import os
import shlex
import subprocess
import sys
import tempfile
from typing import Dict, List, Mapping, Optional, Tuple

from lib_py.agents.base import AGENT_SESSION_ENV
from lib_py.agents.registry import agent_of_row, own_key
from lib_py.update import open_jobs, row_server

SPEC_VERSION = 1
SPEC_REL = os.path.join(".mam", "reset-spec.json")


class SpecError(Exception):
    pass


def _real(p: str) -> str:
    return os.path.realpath(p) if p else ""


def spec_path(workspace: str) -> str:
    return os.path.join(_real(workspace), SPEC_REL)


def _state() -> dict:
    return json.loads(os.environ.get("MAM_STATE_JSON") or "{}") or {}


# --------------------------------------------------------------------------- D2: target selection

def select_targets(workspace: str, state: dict) -> dict:
    """Select running control-root sessions and classify non-targets.

    Returns {"targets": [...], "skipped": [{"name", "reason"}],
             "orchestrators": [...], "errors": [...]}.
    """
    ws_real = _real(workspace)
    targets: List[dict] = []
    skipped: List[dict] = []
    errors: List[str] = []

    for row in state.get("herdr_sessions") or []:
        if not isinstance(row, dict):
            continue
        name = str(row.get("name") or "")
        status = str(row.get("status") or "")
        if status != "running":
            skipped.append({"name": name, "reason": "status %s" % (status or "unknown")})
            continue

        cwd = str((row.get("pane") or {}).get("cwd") or "")
        cwd_real = _real(cwd) if cwd else ""
        if cwd_real != ws_real:
            skipped.append({
                "name": name,
                "reason": "pane.cwd %s is not the control root (a worktree creator?)" % (cwd or "<empty>"),
            })
            continue

        agent = agent_of_row(row)
        role = str(row.get("role") or "").strip()
        if not agent:
            errors.append(
                "running session %s has no recognisable agent type (agent=%r, pane.cmd=%r); "
                "fix its registry row first"
                % (name, row.get("agent"), (row.get("pane") or {}).get("cmd"))
            )
            continue
        if not role:
            errors.append(
                "running session %s (%s) has an empty role; fix its registry row first"
                % (name, agent)
            )
            continue

        key = own_key(agent)
        raw_id = row.get(key)
        old_id = str(raw_id).strip() if isinstance(raw_id, str) and raw_id.strip() else None
        targets.append({
            "name": name,
            "agent": agent,
            "role": role,
            "herdr_server": str(row_server(row)),
            "herdr_workspace": str(row.get("herdr_workspace") or ""),
            "old_id": old_id,
            "old_created_at": str(row.get("herdr_session_created_at") or ""),
        })

    orchestrators = [str(u) for u in (state.get("orchestrator_uuids") or []) if u]
    return {
        "targets": targets,
        "skipped": skipped,
        "orchestrators": orchestrators,
        "errors": errors,
    }


# --------------------------------------------------------------------------- D5: spec file

def write_spec(workspace: str, agents: List[dict]) -> str:
    """Atomically write .mam/reset-spec.json (mkstemp -> fsync -> replace -> dir fsync)."""
    ws_real = _real(workspace)
    mam_dir = os.path.join(ws_real, ".mam")
    os.makedirs(mam_dir, exist_ok=True)
    path = os.path.join(mam_dir, "reset-spec.json")
    now_iso = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    doc = {
        "version": SPEC_VERSION,
        "workspace": ws_real,
        "written_at": now_iso,
        "agents": agents,
    }
    payload = (json.dumps(doc, indent=2) + "\n").encode("utf-8")
    fd, tmp_path = tempfile.mkstemp(prefix=".reset-spec.", suffix=".tmp", dir=mam_dir)
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(payload)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(tmp_path, path)
    except BaseException:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass
        raise
    try:
        dir_fd = os.open(mam_dir, os.O_RDONLY)
        try:
            os.fsync(dir_fd)
        finally:
            os.close(dir_fd)
    except OSError:
        pass
    return path


def read_spec(workspace: str) -> dict:
    """Read and validate .mam/reset-spec.json; raise SpecError if invalid."""
    ws_real = _real(workspace)
    path = spec_path(ws_real)
    try:
        with open(path, encoding="utf-8") as fh:
            doc = json.load(fh)
    except (OSError, ValueError) as exc:
        raise SpecError("spec file %s cannot be read (%s); check it, then remove it by hand" % (path, exc))

    if not isinstance(doc, dict) or doc.get("version") != SPEC_VERSION:
        raise SpecError(
            "spec file %s has unsupported version %r; check it, then remove it by hand"
            % (path, doc.get("version") if isinstance(doc, dict) else None)
        )
    spec_ws = _real(str(doc.get("workspace") or ""))
    if spec_ws != ws_real:
        raise SpecError(
            "spec file %s is for workspace %r, not %r; check it, then remove it by hand"
            % (path, doc.get("workspace"), ws_real)
        )
    agents = doc.get("agents")
    if not isinstance(agents, list):
        raise SpecError("spec file %s has no agents list; check it, then remove it by hand" % path)
    for item in agents:
        if not isinstance(item, dict):
            raise SpecError("spec file %s has a non-object agent entry; check it, then remove it by hand" % path)
        name = item.get("name")
        agent = item.get("agent")
        role = item.get("role")
        if not (isinstance(name, str) and name and isinstance(agent, str) and agent and isinstance(role, str) and role):
            raise SpecError("spec file %s has an incomplete agent entry; check it, then remove it by hand" % path)
        try:
            own_key(agent)
        except KeyError:
            raise SpecError(
                "spec file %s has unknown agent %r; check it, then remove it by hand" % (path, agent)
            )
    return doc


def new_running_rows(workspace: str, state: dict, spec_agents: List[dict]) -> List[str]:
    """Names of running control-root rows not present in spec_agents."""
    known = {str(a.get("name") or "") for a in spec_agents}
    sel = select_targets(workspace, state)
    return [t["name"] for t in sel["targets"] if t["name"] not in known]


# --------------------------------------------------------------------------- D6: progress classification

def classify_entry(entry: dict, row: Optional[dict]) -> Tuple[str, Optional[str], bool]:
    """Classify a spec entry against its current registry row (matched by name only).

    Returns (r_state, new_id, same_instance):
      1. row is None -> ("purged", None, False)
      2. r.herdr_session_created_at == old_created_at -> ("pending", None, True)
      3. not same_instance and r[id] == old_id (non-empty) -> ("pending", cur_id, False)
      4. not same_instance and r[id] is empty -> ("created", None, False)
      5. not same_instance and r[id] is non-empty and != old_id -> ("done", cur_id, False)
    """
    if row is None:
        return "purged", None, False

    key = own_key(entry["agent"])
    cur_created_at = str(row.get("herdr_session_created_at") or "")
    old_created_at = str(entry.get("old_created_at") or "")
    raw_cur = row.get(key)
    cur_id = str(raw_cur).strip() if isinstance(raw_cur, str) and raw_cur.strip() else None
    raw_old = entry.get("old_id")
    old_id = str(raw_old).strip() if isinstance(raw_old, str) and raw_old.strip() else None

    if cur_created_at == old_created_at:
        return "pending", None, True
    if cur_id and old_id and cur_id == old_id:
        return "pending", cur_id, False
    if not cur_id:
        return "created", None, False
    return "done", cur_id, False


# --------------------------------------------------------------------------- D4: precheck

def _ancestor_pids(start_pid: Optional[int] = None) -> set:
    pid = os.getpid() if start_pid is None else start_pid
    try:
        res = subprocess.run(
            ["ps", "-eo", "pid=,ppid="],
            stdin=subprocess.DEVNULL,
            capture_output=True,
            text=True,
            timeout=5,
        )
    except (OSError, subprocess.SubprocessError):
        return set()
    if res.returncode != 0:
        return set()
    ppid_of: Dict[int, int] = {}
    for line in (res.stdout or "").splitlines():
        parts = line.split()
        if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit():
            ppid_of[int(parts[0])] = int(parts[1])
    ancestors = set()
    cur = pid
    while cur and cur > 1 and cur not in ancestors:
        ancestors.add(cur)
        cur = ppid_of.get(cur, 0)
    return ancestors


def _caller_blockers(
    targets: List[dict],
    rows_by_name: Mapping[str, dict],
    env: Mapping[str, str],
    ancestors: Optional[set] = None,
) -> List[str]:
    env_ids = {
        str(env.get(k) or "").strip()
        for k in AGENT_SESSION_ENV
        if str(env.get(k) or "").strip()
    }
    anc = _ancestor_pids() if ancestors is None else ancestors
    out: List[str] = []
    for t in targets:
        name = t["name"]
        row = rows_by_name.get(name)
        r_state, _, _ = classify_entry(t, row)
        if r_state != "pending":
            continue
        old_id = t.get("old_id")
        key = own_key(t["agent"])
        raw_cur = (row or {}).get(key)
        cur_id = str(raw_cur).strip() if isinstance(raw_cur, str) and raw_cur.strip() else None

        matched = False
        if env_ids and ((old_id and old_id in env_ids) or (cur_id and cur_id in env_ids)):
            matched = True
        elif row is not None and anc:
            pane_pid = (row.get("pane") or {}).get("pid")
            child_pid = row.get("child_pid")
            for p in (pane_pid, child_pid):
                if isinstance(p, int) and p > 1 and p in anc:
                    matched = True
                    break
        if matched:
            out.append("you are %s, one of the targets; run the reset from the orchestrator" % name)
    return out


def precheck(
    workspace: str,
    state: dict,
    env: Mapping[str, str],
    ancestors: Optional[set] = None,
) -> List[str]:
    """Return refusal lines for open jobs, unreadable spec, un-recreatable targets, or caller==target."""
    ws_real = _real(workspace)
    blockers: List[str] = []

    spath = spec_path(ws_real)
    targets: List[dict] = []
    if os.path.lexists(spath):
        try:
            spec = read_spec(ws_real)
            targets = spec["agents"]
        except SpecError as exc:
            blockers.append(str(exc))
    else:
        sel = select_targets(ws_real, state)
        blockers.extend(sel["errors"])
        targets = sel["targets"]

    for j in open_jobs(ws_real, state, env):
        if j.get("dir"):
            blockers.append(
                "jobs directory %s cannot be listed; fix its permissions, then try again"
                % j["path"]
            )
        elif j.get("unreadable"):
            blockers.append(
                "job record %s cannot be read; check it, then move or delete it by hand"
                % j["path"]
            )
        else:
            blockers.append(
                "job %s (%s, %s) is open; wait for it to end or close it: python "
                ".agents/skills/multi-agent-mux-delegate-job/scripts/registry.py"
                " --registry-dir %s close --job %s --request cancelled"
                % (
                    j["job_id"],
                    j["agent_session"] or "-",
                    j["status"],
                    j["registry_dir"],
                    j["job_id"],
                )
            )

    rows_by_name = {
        str(s.get("name")): s
        for s in (state.get("herdr_sessions") or [])
        if isinstance(s, dict) and s.get("name")
    }
    blockers.extend(_caller_blockers(targets, rows_by_name, env, ancestors=ancestors))
    return blockers


# --------------------------------------------------------------------------- D1: entries shell lines

def format_entry_line(entry: dict, row: Optional[dict]) -> str:
    r_state, _, _ = classify_entry(entry, row)
    key = own_key(entry["agent"])
    raw_cur = (row or {}).get(key)
    row_id = str(raw_cur).strip() if isinstance(raw_cur, str) and raw_cur.strip() else ""
    pairs = [
        ("R_STATE", r_state),
        ("R_NAME", str(entry.get("name") or "")),
        ("R_AGENT", str(entry.get("agent") or "")),
        ("R_ROLE", str(entry.get("role") or "")),
        ("R_SERVER", str(entry.get("herdr_server") or "default")),
        ("R_LABEL", str(entry.get("herdr_workspace") or "")),
        ("R_OLD_ID", str(entry.get("old_id") or "")),
        ("R_ROW_ID", row_id),
    ]
    return " ".join("%s=%s" % (k, shlex.quote(v)) for k, v in pairs)


# --------------------------------------------------------------------------- D8: result table

def _parse_skipped(items: Optional[List[str]]) -> Dict[str, str]:
    out: Dict[str, str] = {}
    for raw in items or []:
        for token in raw.split(","):
            token = token.strip()
            if not token:
                continue
            if ":" in token:
                name, kind = token.split(":", 1)
                out[name] = kind
            else:
                out[token] = "purge_refused"
    return out


def _parse_failed(raw: Optional[str]) -> Optional[Tuple[str, str, str]]:
    if not raw:
        return None
    parts = raw.split(":", 2)
    if len(parts) == 3:
        return parts[0], parts[1], parts[2]
    if len(parts) == 2:
        return parts[0], parts[1], "1"
    return parts[0], "stop", "1"


def build_table(
    agents: List[dict],
    state: dict,
    skipped_items: Optional[List[str]] = None,
    failed_spec: Optional[str] = None,
) -> Tuple[str, int]:
    """Format the D8 result table and compute the D9 exit code (0, 1, 4, or 5)."""
    rows_by_name = {
        str(s.get("name")): s
        for s in (state.get("herdr_sessions") or [])
        if isinstance(s, dict) and s.get("name")
    }
    skipped_map = _parse_skipped(skipped_items)
    failed = _parse_failed(failed_spec)
    failed_name = failed[0] if failed else None
    after_failed = False

    headers = ("NAME", "AGENT", "OLD_ID", "NEW_ID", "RECORDED", "NEW≠OLD", "RESULT")
    table_rows: List[Tuple[str, str, str, str, str, str, str]] = []

    has_skip_or_incomplete = bool(skipped_map)
    has_created = False

    for entry in agents:
        name = str(entry.get("name") or "")
        agent = str(entry.get("agent") or "")
        raw_old = entry.get("old_id")
        old_id = str(raw_old).strip() if isinstance(raw_old, str) and raw_old.strip() else None
        old_col = old_id if old_id else "-"

        row = rows_by_name.get(name)
        r_state, new_id, _ = classify_entry(entry, row)
        new_col = new_id if new_id else "-"
        recorded_col = "yes" if new_id else "no"
        if not new_id:
            diff_col = "-"
        elif not old_id or new_id != old_id:
            diff_col = "yes"
        else:
            diff_col = "no"

        if failed and name == failed_name:
            result_col = "failed at %s (rc %s)" % (failed[1], failed[2])
            after_failed = True
            has_skip_or_incomplete = True
        elif after_failed and r_state in ("pending", "purged"):
            result_col = "not reached"
            has_skip_or_incomplete = True
        elif name in skipped_map:
            kind = skipped_map[name]
            if kind == "no_id":
                result_col = "skipped: no conversation id recorded"
            else:
                result_col = "skipped: purge refused"
            has_skip_or_incomplete = True
        elif r_state == "done":
            result_col = "done"
        elif r_state == "created":
            result_col = "created; id not recorded yet"
            has_created = True
        elif r_state == "purged":
            result_col = "purged"
            has_skip_or_incomplete = True
        else:
            result_col = "pending"
            has_skip_or_incomplete = True

        table_rows.append((name, agent, old_col, new_col, recorded_col, diff_col, result_col))

    widths = [len(h) for h in headers]
    for r in table_rows:
        for i in range(len(headers) - 1):
            if len(r[i]) > widths[i]:
                widths[i] = len(r[i])

    def _fmt(cols: Tuple[str, ...]) -> str:
        parts = [cols[i].ljust(widths[i]) for i in range(len(cols) - 1)]
        parts.append(cols[-1])
        return "  ".join(parts)

    lines = [_fmt(headers)]
    for r in table_rows:
        lines.append(_fmt(r))

    if failed:
        rc = 1
    elif has_skip_or_incomplete:
        rc = 4
    elif has_created:
        rc = 5
    else:
        rc = 0
    return "\n".join(lines), rc


# --------------------------------------------------------------------------- D10: dry-run

def format_dry_run(
    workspace: str,
    state: dict,
    env: Mapping[str, str],
    extra_refusals: Optional[List[str]] = None,
) -> Tuple[str, int]:
    ws_real = _real(workspace)
    lines: List[str] = []
    refusals: List[str] = list(extra_refusals or [])
    refusals.extend(precheck(ws_real, state, env))

    rows_by_name = {
        str(s.get("name")): s
        for s in (state.get("herdr_sessions") or [])
        if isinstance(s, dict) and s.get("name")
    }
    sel = select_targets(ws_real, state)
    spath = spec_path(ws_real)
    action_map = {
        "pending": "stop + purge + create",
        "purged": "create",
        "created": "none (waiting for id)",
        "done": "none (done)",
    }

    if os.path.lexists(spath):
        try:
            spec = read_spec(ws_real)
        except SpecError:
            spec = None
        if spec is not None:
            lines.append("spec file: %s" % spath)
            for e in spec["agents"]:
                r_state, _, _ = classify_entry(e, rows_by_name.get(e["name"]))
                lines.append(
                    "target: %s  agent=%s  role=%s  herdr_server=%s  herdr_workspace=%s  "
                    "old_id=%s  old_created_at=%s  state=%s  action=%s"
                    % (
                        e["name"],
                        e["agent"],
                        e["role"],
                        e.get("herdr_server") or "default",
                        e.get("herdr_workspace") or "",
                        e.get("old_id") or "-",
                        e.get("old_created_at") or "-",
                        r_state,
                        action_map.get(r_state, r_state),
                    )
                )
            extra_running = new_running_rows(ws_real, state, spec["agents"])
            if extra_running:
                lines.append(
                    "NOTE: running control-root session(s) not in reset-spec.json "
                    "(not targeted by this rerun): %s" % ", ".join(extra_running)
                )
    else:
        if sel["targets"]:
            for t in sel["targets"]:
                lines.append(
                    "target: %s  agent=%s  role=%s  herdr_server=%s  herdr_workspace=%s  "
                    "old_id=%s  old_created_at=%s"
                    % (
                        t["name"],
                        t["agent"],
                        t["role"],
                        t["herdr_server"],
                        t["herdr_workspace"],
                        t["old_id"] or "-",
                        t["old_created_at"] or "-",
                    )
                )
        else:
            lines.append("targets: (none)")

    for sk in sel["skipped"]:
        lines.append("skip: %s (%s)" % (sk["name"], sk["reason"]))
    for u in sel["orchestrators"]:
        lines.append("orchestrator %s (not a row; never reset)" % u)

    if refusals:
        lines.append("would refuse:")
        for r in refusals:
            lines.append("  - %s" % r)
        return "\n".join(lines), 1
    return "\n".join(lines), 0


# --------------------------------------------------------------------------- CLI

def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(prog="lib_py.reset")
    sub = p.add_subparsers(dest="cmd", required=True)

    s = sub.add_parser("plan")
    s.add_argument("--workspace", required=True)

    s = sub.add_parser("precheck")
    s.add_argument("--workspace", required=True)

    s = sub.add_parser("write-spec")
    s.add_argument("--workspace", required=True)

    s = sub.add_parser("entries")
    s.add_argument("--workspace", required=True)
    s.add_argument("--session", default="")
    s.add_argument("--warn-new", action="store_true")

    s = sub.add_parser("table")
    s.add_argument("--workspace", required=True)
    s.add_argument("--skipped", action="append", default=[])
    s.add_argument("--failed", default="")

    s = sub.add_parser("dry-run")
    s.add_argument("--workspace", required=True)
    s.add_argument("--refusal", action="append", default=[])

    a = p.parse_args(argv)
    ws = _real(a.workspace)

    if a.cmd == "write-spec":
        try:
            raw = json.load(sys.stdin)
        except ValueError as exc:
            print("ERROR: invalid targets JSON on stdin (%s)" % exc, file=sys.stderr)
            return 1
        agents = raw.get("targets") if isinstance(raw, dict) else raw
        if not isinstance(agents, list):
            print("ERROR: targets JSON must be a list", file=sys.stderr)
            return 1
        try:
            write_spec(ws, agents)
        except OSError as exc:
            print("ERROR: cannot write %s (%s)" % (spec_path(ws), exc), file=sys.stderr)
            return 1
        return 0

    try:
        state = _state()
    except ValueError as exc:
        if a.cmd == "precheck":
            print("registry: cannot parse MAM_STATE_JSON (%s)" % exc)
            return 1
        print("ERROR: cannot parse MAM_STATE_JSON (%s)" % exc, file=sys.stderr)
        return 1

    if a.cmd == "plan":
        sel = select_targets(ws, state)
        print(json.dumps(sel))
        return 1 if sel["errors"] else 0

    if a.cmd == "precheck":
        lines = precheck(ws, state, os.environ)
        for line in lines:
            print(line)
        return 1 if lines else 0

    if a.cmd == "entries":
        try:
            spec = read_spec(ws)
        except SpecError as exc:
            print("ERROR: %s" % exc, file=sys.stderr)
            return 1
        if a.warn_new:
            extra = new_running_rows(ws, state, spec["agents"])
            if extra:
                print(
                    "NOTE: running control-root session(s) not in %s (not targeted by this rerun): %s"
                    % (spec_path(ws), ", ".join(extra)),
                    file=sys.stderr,
                )
        rows_by_name = {
            str(s.get("name")): s
            for s in (state.get("herdr_sessions") or [])
            if isinstance(s, dict) and s.get("name")
        }
        for e in spec["agents"]:
            if a.session and e.get("name") != a.session:
                continue
            print(format_entry_line(e, rows_by_name.get(e["name"])))
        return 0

    if a.cmd == "table":
        try:
            spec = read_spec(ws)
        except SpecError as exc:
            print("ERROR: %s" % exc, file=sys.stderr)
            return 1
        out, rc = build_table(
            spec["agents"],
            state,
            skipped_items=a.skipped,
            failed_spec=a.failed or None,
        )
        print(out)
        return rc

    if a.cmd == "dry-run":
        out, rc = format_dry_run(ws, state, os.environ, extra_refusals=a.refusal)
        print(out)
        return rc

    return 2


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