"""Same-cwd creation/resume ordering for post-discovered agents (Issue #63).

While a post-discovered agent (agy, codex, hermes, opencode) in a workspace is
 still being created (marker held by a live create/resume process) or is running
without a recorded conversation id (live pane pid), a second create or resume of
the same agent type in that workspace waits until the first records its id, up
to MAM_ID_ORDER_WAIT_S (default 600 s; 0 = do not wait). A marker whose owner
died before the row was registered keeps blocking while its herdr session is
alive (the pane may still discover the workspace id).
"""
from __future__ import annotations

import argparse
import fcntl
import hashlib
import json
import os
import subprocess
import sys
import tempfile
import time
from typing import Any, Callable, Dict, List, Mapping, Optional

from lib_py import ints
from lib_py.agents.autopin import claimed_elsewhere
from lib_py.agents.base import POST_DISCOVERED
from lib_py.agents.registry import agent_of_row, get_adapter
from lib_py.agents.session_ids import canonical_cwd, claiming_rows, own_id
from lib_py.state import load_state_json

DEFAULT_WAIT_S = 600


def wait_s(env: Mapping[str, str] = os.environ) -> int:
    raw = env.get("MAM_ID_ORDER_WAIT_S")
    if raw is None:
        return DEFAULT_WAIT_S
    n = ints.positive_int(raw, 0)
    return DEFAULT_WAIT_S if n is None else n


def pid_alive(pid: int) -> bool:
    if not isinstance(pid, int) or pid <= 0:
        return False
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    except OSError:
        return False


def herdr_session_alive(server: str, session: str) -> bool:
    """lib.sh `_herdr has-session` on the marker's herdr server; an unanswered
    probe counts as alive (the wait stays bounded by MAM_ID_ORDER_WAIT_S)."""
    skills_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    env = dict(os.environ)
    if server:
        env["HERDR_SESSION_NAME"] = server
    script = 'source "$1" && _herdr has-session -t "$2"'
    try:
        return subprocess.run(
            ["bash", "-c", script, "_", os.path.join(skills_dir, "lib.sh"), session],
            env=env, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL, timeout=30,
        ).returncode == 0
    except (OSError, subprocess.SubprocessError):
        return True


def order_dir(cwd: str, agent: str, env: Mapping[str, str] = os.environ,
              home: Optional[str] = None) -> str:
    h = home or env.get("HOME") or env.get("HOME_DIR") or os.path.expanduser("~")
    real_cwd = canonical_cwd(cwd or ".")
    digest = hashlib.sha256(real_cwd.encode("utf-8")).hexdigest()[:16]
    return os.path.join(h, ".local", "state", "mam", "id-order", f"{digest}-{agent}")


def _parse_pos_int(val: Any) -> Optional[int]:
    if isinstance(val, bool):
        return None
    if isinstance(val, int):
        return val if val > 0 else None
    s = str(val or "").strip()
    return int(s) if s.isdigit() and int(s) > 0 else None


def blockers(state: Dict[str, Any], cwd: str, agent: str, me: str,
             markers: List[Dict[str, Any]],
             is_alive: Callable[[int], bool] = pid_alive) -> List[str]:
    target_cwd = canonical_cwd(cwd)
    found: List[str] = []
    rows_by_name: Dict[str, Dict[str, Any]] = {}

    for s in state.get("herdr_sessions") or []:
        if not isinstance(s, dict):
            continue
        rname = str(s.get("name") or "")
        if rname:
            rows_by_name[rname] = s
        if not rname or rname == me:
            continue
        s_cwd = (s.get("pane") or {}).get("cwd") or ""
        if not s_cwd or canonical_cwd(s_cwd) != target_cwd:
            continue
        if agent_of_row(s) != agent or s.get("status") != "running":
            continue
        if own_id(s, agent):
            continue
        pane_pid = _parse_pos_int((s.get("pane") or {}).get("pid"))
        if pane_pid is None or is_alive(pane_pid):
            if rname not in found:
                found.append(rname)

    for m in markers or []:
        if not isinstance(m, dict):
            continue
        m_sess = str(m.get("session") or "")
        if not m_sess or m_sess in found:
            continue
        m_cwd = str(m.get("cwd") or "")
        if m_cwd and canonical_cwd(m_cwd) != target_cwd:
            continue
        m_agent = str(m.get("agent") or "")
        if m_agent and m_agent != agent:
            continue
        row = rows_by_name.get(m_sess)
        if row is None:
            found.append(m_sess)
            continue
        status = str(row.get("status") or "")
        cur_own = own_id(row, agent)
        row_epoch = int(row.get("herdr_session_epoch") or 0)
        end_epoch = int(row.get("stopped_at_epoch") or row.get("terminated_at_epoch") or 0)
        m_at = int(m.get("at") or 0)
        prev_own = str(m.get("prev_own_id") or "")
        prev_epoch = int(m.get("prev_epoch") or 0)
        prev_end_epoch = int(m.get("prev_end_epoch") or 0)

        pinned_for_attempt = bool(cur_own) and (
            cur_own != prev_own
            or (status == "running" and row_epoch > prev_epoch)
        )
        if pinned_for_attempt:
            continue
        if status == "running":
            found.append(m_sess)
        else:
            ended_after_attempt = (
                end_epoch > max(m_at, prev_end_epoch)
                and (prev_epoch == 0 or row_epoch > prev_epoch)
            )
            if not ended_after_attempt:
                found.append(m_sess)

    return found


def _registered_for_attempt(m: Dict[str, Any], row: Optional[Dict[str, Any]]) -> bool:
    if not isinstance(row, dict) or row.get("status") != "running":
        return False
    return int(row.get("herdr_session_epoch") or 0) > int(m.get("prev_epoch") or 0)


def _read_and_prune_markers(dpath: str, me: str, my_pid: int, state: Dict[str, Any],
                            is_alive: Callable[[int], bool] = pid_alive,
                            pane_alive: Callable[[str, str], bool] = herdr_session_alive,
                            ) -> List[Dict[str, Any]]:
    """Markers that still count. A live owner keeps its marker (another live
    attempt on `me` included). A dead owner's marker goes once its row is
    registered for the attempt (the row decides from then on) or its herdr
    session is gone; until then it is kept as an orphan."""
    rows_by_name = {
        str(s.get("name") or ""): s
        for s in (state.get("herdr_sessions") or []) if isinstance(s, dict)
    }
    out: List[Dict[str, Any]] = []
    try:
        names = sorted(os.listdir(dpath))
    except OSError:
        return out
    for fname in names:
        if not fname.endswith(".json") or fname.startswith("."):
            continue
        fpath = os.path.join(dpath, fname)
        try:
            with open(fpath, encoding="utf-8") as f:
                data = json.load(f)
        except (OSError, ValueError):
            try:
                os.remove(fpath)
            except OSError:
                pass
            continue
        if not isinstance(data, dict):
            continue
        owner_pid = _parse_pos_int(data.get("owner_pid"))
        sess = str(data.get("session") or "")
        if owner_pid is None or owner_pid == my_pid:
            drop = True
        elif is_alive(owner_pid):
            drop = False
        elif sess == me or not sess:
            drop = True
        elif _registered_for_attempt(data, rows_by_name.get(sess)):
            drop = True
        else:
            drop = not pane_alive(str(data.get("herdr_server") or ""), sess)
            data["_orphan"] = True
        if drop:
            try:
                os.remove(fpath)
            except OSError:
                pass
            continue
        out.append(data)
    return out


def write_marker(dpath: str, session: str, owner_pid: int, cwd: str, agent: str,
                 prev_row: Optional[Dict[str, Any]] = None, herdr_server: str = "") -> str:
    os.makedirs(dpath, exist_ok=True)
    marker_path = os.path.join(dpath, f"{session}.json")
    payload: Dict[str, Any] = {
        "session": session,
        "owner_pid": int(owner_pid),
        "cwd": canonical_cwd(cwd),
        "agent": agent,
        "herdr_server": herdr_server,
        "at": int(time.time()),
    }
    if isinstance(prev_row, dict):
        payload["prev_own_id"] = own_id(prev_row, agent)
        payload["prev_epoch"] = int(prev_row.get("herdr_session_epoch") or 0)
        payload["prev_end_epoch"] = int(
            prev_row.get("stopped_at_epoch") or prev_row.get("terminated_at_epoch") or 0
        )
    fd, tmp = tempfile.mkstemp(dir=dpath, prefix=".marker.", suffix=".tmp")
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        json.dump(payload, f)
    os.replace(tmp, marker_path)
    return marker_path


def wait(cwd: str, agent: str, session: str, owner_pid: int,
         uuid: str = "", yaml_path: Optional[str] = None,
         home: Optional[str] = None, wait_s_val: Optional[int] = None,
         poll: float = 0.25,
         is_alive: Callable[[int], bool] = pid_alive,
         herdr_server: str = "",
         pane_alive: Optional[Callable[[str, str], bool]] = None) -> int:
    adapter = get_adapter(agent)
    if adapter is None or adapter.id_lifecycle != POST_DISCOVERED:
        return 0

    cap = wait_s(os.environ) if wait_s_val is None else int(wait_s_val)
    dpath = order_dir(cwd, agent, os.environ, home=home)
    os.makedirs(dpath, exist_ok=True)
    lock_file = os.path.join(dpath, "lock")
    deadline = time.monotonic() + cap
    announced = False
    row_ref = {"name": session}
    probe = pane_alive or herdr_session_alive

    while True:
        with open(lock_file, "a", encoding="utf-8") as lf:
            fcntl.flock(lf, fcntl.LOCK_EX)
            state = load_state_json(yaml_path)
            if uuid:
                claimed = set(claimed_elsewhere(row_ref, state))
                if uuid in claimed:
                    others = ", ".join(claiming_rows(row_ref, uuid, state))
                    print(
                        f"ERROR: {session}'s recorded id {uuid} is also recorded by {others}; "
                        f"refusing to resume on a conversation another session claims (Issue #63).",
                        file=sys.stderr, flush=True,
                    )
                    return 1
            markers = _read_and_prune_markers(dpath, session, owner_pid, state,
                                              is_alive=is_alive, pane_alive=probe)
            blks = blockers(state, cwd, agent, session, markers, is_alive=is_alive)
            if not blks:
                prev_row = next(
                    (
                        s for s in (state.get("herdr_sessions") or [])
                        if isinstance(s, dict) and s.get("name") == session
                    ),
                    None,
                )
                already_recorded = bool(uuid) and own_id(prev_row, agent) == uuid
                if not already_recorded and owner_pid and owner_pid > 0:
                    write_marker(dpath, session, owner_pid, cwd, agent, prev_row=prev_row,
                                 herdr_server=herdr_server)
                return 0

        names = ", ".join(blks)
        if time.monotonic() >= deadline:
            print(
                f"ERROR: {names} in {cwd} still has no recorded {agent} id after {cap}s; "
                f"not starting {session} (Issue #63). Wait for its first turn to finish, "
                f"or stop it with multi-agent-mux-stop.",
                file=sys.stderr, flush=True,
            )
            for m in markers:
                if m.get("_orphan") and m.get("session") in blks:
                    srv = m.get("herdr_server") or "default"
                    print(
                        f"ERROR: {m['session']}'s create/resume exited before recording it, "
                        f"but its herdr session is alive; close it with: "
                        f"HERDR_SESSION_NAME={srv} herdr kill-session -t {m['session']}",
                        file=sys.stderr, flush=True,
                    )
            return 1
        if not announced:
            print(
                f"id-order: waiting for {names} ({agent}, no recorded id yet) in {cwd} "
                f"- up to {cap}s (MAM_ID_ORDER_WAIT_S)",
                file=sys.stderr, flush=True,
            )
            announced = True
        remaining = deadline - time.monotonic()
        if remaining > 0:
            time.sleep(min(poll, remaining))


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(prog="python -m lib_py.agents.id_order")
    sub = p.add_subparsers(dest="cmd", required=True)
    s_wait = sub.add_parser("wait")
    s_wait.add_argument("--cwd", required=True)
    s_wait.add_argument("--agent", required=True)
    s_wait.add_argument("--session", required=True)
    s_wait.add_argument("--owner-pid", type=int, default=0)
    s_wait.add_argument("--uuid", default="")
    s_wait.add_argument("--herdr-server", default="")
    a = p.parse_args(argv)
    return wait(
        cwd=a.cwd,
        agent=a.agent,
        session=a.session,
        owner_pid=a.owner_pid,
        uuid=a.uuid or "",
        herdr_server=a.herdr_server or "",
    )


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