"""Issue #50 follow-up: one pane claim at a time per herdr server.

The herdr shim's `new-session` chooses a pane - an idle one in a workspace on
the same cwd (BUG-4), a split of one, or the root pane of a fresh workspace -
and then claims it with `agent start`. Two shims on one server (several
`create` at once into one folder) could choose the same idle pane in between,
and the later `agent start` failed with agent_pane_busy. The shim now holds
one lock per server from the choice to the claim.

  pane_lock.py path --session S   the lock file of herdr session S, under the
                                  HOME whose socket the session uses. S empty
                                  (or `default`) is the shim's ambient target:
                                  herdr then answers on the server of the pane
                                  this runs in (HERDR_SESSION / HERDR_SOCKET_PATH,
                                  as lib.sh _herdr_caller_server reads them), or
                                  the default server outside any pane - the
                                  lock is that server's (review e8e471ff R1)
  pane_lock.py acquire --fd N     flock(LOCK_EX) on the inherited fd N; prints
                                  `locked`, or `timeout` once MAM_PANE_CLAIM_WAIT_S
                                  (default 120, minimum 1) has passed

The lock is a BSD flock, which belongs to the open file the shim holds on fd 9:
it outlives this helper and ends with the shim's process, however that ends
(exit, kill -9), so a dead holder leaves nothing to steal. flock(1) is not on
macOS; fcntl.flock is. Stdlib only, Python 3.9.
"""
from __future__ import annotations

import argparse
import fcntl
import os
import sys
import time
from typing import List, Mapping, Optional

from lib_py import ints

DEFAULT_WAIT_S = 120


def caller_server(env: Mapping[str, str]) -> str:
    """The server whose pane this process runs in, by the pane's exports:
    HERDR_SESSION when present, else HERDR_SOCKET_PATH's layout
    ($HOME/.config/herdr/herdr.sock -> default, .../sessions/<name>/herdr.sock
    -> name); "" outside a pane or for an unrecognised layout."""
    if env.get("HERDR_SESSION"):
        return env["HERDR_SESSION"]
    sock = env.get("HERDR_SOCKET_PATH") or ""
    if sock.endswith("/.config/herdr/herdr.sock"):
        return "default"
    if sock.endswith("/herdr.sock") and "/.config/herdr/sessions/" in sock:
        return sock[:-len("/herdr.sock")].rsplit("/", 1)[-1]
    return ""


def lock_path(session: Optional[str], env: Mapping[str, str]) -> str:
    home = env.get("HOME") or env.get("HOME_DIR") or os.path.expanduser("~")
    if session in (None, "", "default"):
        session = caller_server(env) or "default"
    return os.path.join(home, ".local", "state", "mam", "pane-lock", session + ".lock")


def wait_s(env: Mapping[str, str]) -> int:
    return ints.positive_int(env.get("MAM_PANE_CLAIM_WAIT_S"), 1) or DEFAULT_WAIT_S


def acquire(fd: int, wait: float) -> str:
    """`locked`, or `timeout` when the lock was not free within `wait` s. The
    deadline is on the monotonic clock; a step of the wall clock does not
    stretch or cut it."""
    deadline = time.monotonic() + wait
    while True:
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            return "locked"
        except OSError:
            if time.monotonic() >= deadline:
                return "timeout"
            time.sleep(0.2)


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(prog="pane_lock.py")
    sub = p.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("path")
    s.add_argument("--session", default="")
    s = sub.add_parser("acquire")
    s.add_argument("--fd", type=int, required=True)
    a = p.parse_args(argv)
    if a.cmd == "path":
        print(lock_path(a.session, os.environ))
    else:
        print(acquire(a.fd, wait_s(os.environ)))
    return 0


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