"""Root-owned tmux client mediator for cross-UID seat control."""

from __future__ import annotations

import argparse
import json
import os
import pwd
import stat
import subprocess
import tempfile
from dataclasses import dataclass

from seat_common import CONTROL_RE, SEAT_ID_RE, die as _common_die
from seat_implementer_identity import (
    expected_socket_for_seat,
    resolve_implementer_for_seat,
)
from seat_scope_entry import (
    exit_receipt_path,
    load_session_manifest,
    parse_proc_start_time,
    publish_exit_receipt,
    resolve_target,
    session_manifest_path,
)

INSTALLED_BIN = "/usr/local/bin/overdeck-seat-tmux-mediator"
REAP_CLOSE_HELPER = os.path.join(os.path.dirname(os.path.abspath(__file__)), "agent-session-reap-close.py")
ALLOWED_SUBCOMMANDS = frozenset({
    "attach-session",
    "list-panes",
    "has-session",
    "kill-server",
    "lifecycle",
    "operator-stop",
    "receipt",
})


def die(msg: str, code: int = 2) -> None:
    _common_die(f"overdeck-seat-tmux-mediator: {msg}", code)


def stop_intent_path(manifest_path: str) -> str:
    return os.path.join(os.path.dirname(manifest_path), "stop-intent.json")


def write_stop_intent(manifest_path: str, manifest: dict[str, object]) -> None:
    path = stop_intent_path(manifest_path)
    parent = os.path.dirname(path)
    fd, temporary = tempfile.mkstemp(prefix=".stop-intent.", dir=parent)
    try:
        os.fchmod(fd, 0o600)
        payload = {
            "schema": 1, "seatId": manifest["seatId"], "generation": manifest["generation"],
            "serverPid": manifest["serverPid"], "serverStartTime": manifest["serverStartTime"],
            "cgroup": manifest["cgroup"],
        }
        with os.fdopen(fd, "w", encoding="utf-8", closefd=False) as output:
            json.dump(payload, output, separators=(",", ":"))
            output.write("\n")
            output.flush()
            os.fsync(fd)
        os.replace(temporary, path)
        temporary = ""
    finally:
        os.close(fd)
        if temporary:
            try:
                os.unlink(temporary)
            except FileNotFoundError:
                pass


def load_stop_intent(manifest_path: str) -> dict[str, object] | None:
    try:
        with open(stop_intent_path(manifest_path), encoding="utf-8") as source:
            value = json.load(source)
    except FileNotFoundError:
        return None
    except (OSError, json.JSONDecodeError):
        die("stop-intent-invalid", 1)
    if not isinstance(value, dict):
        die("stop-intent-invalid", 1)
    return value


def remove_stop_intent(manifest_path: str) -> None:
    try:
        os.unlink(stop_intent_path(manifest_path))
    except FileNotFoundError:
        pass


def verify_stopped_identity(manifest: dict[str, object], proc_root: str = "/proc") -> None:
    pid = int(manifest["serverPid"])
    try:
        with open(os.path.join(proc_root, str(pid), "stat"), encoding="utf-8") as source:
            current_start = parse_proc_start_time(source.read())
    except (FileNotFoundError, ProcessLookupError):
        return
    if current_start == int(manifest["serverStartTime"]):
        die("operator-stop-server-still-live", 1)


def validate_text(name: str, value: str) -> None:
    if not value or CONTROL_RE.search(value):
        die(f"invalid {name}")


@dataclass
class SocketAncestryBinding:
    socket: str
    parents: list[tuple[str, int, tuple[int, int]]]
    socket_identity: tuple[int, int]

    def close(self) -> None:
        for _path, fd, _identity in reversed(self.parents):
            os.close(fd)
        self.parents.clear()


def bind_socket_ancestry(socket: str) -> SocketAncestryBinding:
    parent = os.path.dirname(socket)
    components = [part for part in parent.split(os.sep) if part]
    current_path = os.sep
    fd = os.open(os.sep, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
    parents = [(os.sep, fd, (os.fstat(fd).st_dev, os.fstat(fd).st_ino))]
    try:
        for component in components:
            fd = os.open(component, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=fd)
            current_path = os.path.join(current_path, component)
            st = os.fstat(fd)
            parents.append((current_path, fd, (st.st_dev, st.st_ino)))
        socket_stat = os.stat(os.path.basename(socket), dir_fd=parents[-1][1], follow_symlinks=False)
        if not stat.S_ISSOCK(socket_stat.st_mode):
            die("socket-leaf-invalid", 1)
        return SocketAncestryBinding(socket, parents, (socket_stat.st_dev, socket_stat.st_ino))
    except OSError:
        for _path, held, _identity in reversed(parents):
            os.close(held)
        die("socket-ancestry-invalid", 1)


def revalidate_socket_ancestry(binding: SocketAncestryBinding) -> None:
    for path, fd, identity in binding.parents:
        held = os.fstat(fd)
        try:
            current = os.stat(path, follow_symlinks=False)
        except OSError:
            die("socket-ancestry-changed", 1)
        if (not stat.S_ISDIR(current.st_mode)
                or (held.st_dev, held.st_ino) != identity
                or (current.st_dev, current.st_ino) != identity):
            die("socket-ancestry-changed", 1)
    try:
        current_socket = os.stat(
            os.path.basename(binding.socket), dir_fd=binding.parents[-1][1], follow_symlinks=False
        )
    except OSError:
        die("socket-leaf-changed", 1)
    if (not stat.S_ISSOCK(current_socket.st_mode)
            or (current_socket.st_dev, current_socket.st_ino) != binding.socket_identity):
        die("socket-leaf-changed", 1)


def validate_socket(socket: str, operator_home: str, seat_id: str) -> None:
    want = expected_socket_for_seat(operator_home, seat_id)
    if socket != want:
        die("socket-mismatch", 1)
    if not socket.startswith("/") or "//" in socket:
        die("invalid socket", 1)
    if not os.path.exists(os.path.dirname(socket)):
        die("socket-parent-missing", 1)


def verify_live_identity(manifest: dict[str, object], socket: str, proc_root: str = "/proc") -> None:
    required = {"socketDev", "socketIno", "serverPid", "serverStartTime", "cgroup"}
    if not required <= set(manifest):
        die("manifest-identity-missing", 1)
    try:
        socket_stat = os.stat(socket, follow_symlinks=False)
        pid = int(manifest["serverPid"])
        stat_text = open(os.path.join(proc_root, str(pid), "stat"), encoding="utf-8").read()
        cgroup_lines = open(os.path.join(proc_root, str(pid), "cgroup"), encoding="utf-8").read().splitlines()
    except (OSError, TypeError, ValueError):
        die("live-identity-unavailable", 1)
    if not stat.S_ISSOCK(socket_stat.st_mode):
        die("socket-type-mismatch", 1)
    if (socket_stat.st_dev, socket_stat.st_ino) != (manifest["socketDev"], manifest["socketIno"]):
        die("socket-identity-mismatch", 1)
    if parse_proc_start_time(stat_text) != manifest["serverStartTime"]:
        die("server-start-mismatch", 1)
    cgroups = [line.split(":", 2)[2] for line in cgroup_lines if line.startswith("0::")]
    if cgroups != [manifest["cgroup"]]:
        die("server-cgroup-mismatch", 1)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="overdeck-seat-tmux-mediator", add_help=True)
    parser.add_argument("--seat-id", required=True)
    parser.add_argument("--socket", required=True)
    parser.add_argument("--generation", required=True)
    parser.add_argument("tmux_args", nargs=argparse.REMAINDER)
    return parser


def validate_tmux_args(args: list[str]) -> list[str]:
    if not args:
        die("tmux-args-missing", 1)
    if args[0] == "--":
        args = args[1:]
    if not args:
        die("tmux-args-missing", 1)
    sub = args[0]
    if sub not in ALLOWED_SUBCOMMANDS:
        die(f"tmux-subcommand-forbidden:{sub}", 1)
    if "-S" in args:
        die("tmux-socket-flag-forbidden", 1)
    if "-L" in args:
        die("tmux-server-flag-forbidden", 1)
    return args


def exec_tmux_attach(implementer_user: str, binding: SocketAncestryBinding,
                     tmux_args: list[str]) -> None:
    parent_fd = binding.parents[-1][1]
    socket = f"/proc/self/fd/{parent_fd}/{os.path.basename(binding.socket)}"
    try:
        account = pwd.getpwnam(implementer_user)
    except KeyError:
        die("implementer-user-missing", 1)
    os.set_inheritable(parent_fd, True)
    os.setgroups([])
    os.setgid(account.pw_gid)
    os.setuid(account.pw_uid)
    os.execvp("tmux", ["tmux", "-S", socket, *tmux_args])


def run_tmux_bound(implementer_user: str, binding: SocketAncestryBinding,
                   tmux_args: list[str]) -> subprocess.CompletedProcess[str]:
    parent_fd = binding.parents[-1][1]
    socket = f"/proc/self/fd/{parent_fd}/{os.path.basename(binding.socket)}"
    try:
        account = pwd.getpwnam(implementer_user)
    except KeyError:
        die("implementer-user-missing", 1)

    def drop_identity() -> None:
        os.setgroups([])
        os.setgid(account.pw_gid)
        os.setuid(account.pw_uid)

    return subprocess.run(
        ["tmux", "-S", socket, *tmux_args], text=True, capture_output=True,
        pass_fds=(parent_fd,), preexec_fn=drop_identity,
    )


def main(argv: list[str] | None = None) -> None:
    parser = build_parser()
    try:
        ns = parser.parse_args(argv)
    except SystemExit:
        die("unknown arg", 2)

    if not SEAT_ID_RE.match(ns.seat_id):
        die("invalid seat id")

    if os.geteuid() != 0:
        die("requires root", 1)

    operator_user, _operator_uid, _operator_gid, operator_home = resolve_target()
    validate_socket(ns.socket, operator_home, ns.seat_id)
    binding = bind_socket_ancestry(ns.socket)
    state_dir = os.path.dirname(ns.socket)
    manifest_path = session_manifest_path(ns.seat_id)
    manifest = load_session_manifest(manifest_path, ns.seat_id, ns.generation or "")
    if manifest["socket"] != ns.socket:
        die("manifest-socket-mismatch", 1)
    implementer_user, uid, _gid = resolve_implementer_for_seat(ns.seat_id, operator_user)
    if manifest["implementerUid"] != uid:
        die("manifest-uid-mismatch", 1)
    tmux_args = list(ns.tmux_args)
    if tmux_args == ["operator-stop"]:
        pending = load_stop_intent(manifest_path)
        if pending is not None:
            expected = {key: manifest[key] for key in (
                "seatId", "generation", "serverPid", "serverStartTime", "cgroup"
            )}
            if {key: pending.get(key) for key in expected} != expected:
                die("stop-intent-identity-mismatch", 1)
            verify_stopped_identity(manifest)
            publish_exit_receipt(
                manifest_path, exit_receipt_path(ns.seat_id), ns.seat_id,
                manifest["generation"], terminal_reason="operator-stop",
            )
            remove_stop_intent(manifest_path)
            binding.close()
            return
    verify_live_identity(manifest, ns.socket)
    if tmux_args and tmux_args[0] == "reap-close":
        identity = tmux_args[1:]
        if len(identity) != 9:
            die("reap-close-identity-invalid", 1)
        revalidate_socket_ancestry(binding)
        proc = subprocess.run(
            [REAP_CLOSE_HELPER, ns.socket, *identity],
            text=True,
            capture_output=True,
        )
        if proc.returncode != 0 or proc.stdout != "REAPED\n":
            die("reap-close-kept", 1)
        binding.close()
        print("REAPED")
        return
    tmux_args = validate_tmux_args(tmux_args)
    if tmux_args[0] == "lifecycle":
        if len(tmux_args) != 1:
            die("lifecycle-args", 1)
        revalidate_socket_ancestry(binding)
        waited = run_tmux_bound(implementer_user, binding,
                                ["wait-for", f"agent-exit-{manifest['generation']}"])
        if waited.returncode != 0:
            die("lifecycle-wait-failed", 1)
        verify_live_identity(manifest, ns.socket)
        revalidate_socket_ancestry(binding)
        pane = run_tmux_bound(
            implementer_user, binding,
            ["list-panes", "-t", "main", "-F", "#{pane_dead} #{pane_dead_status}"],
        )
        fields = pane.stdout.strip().split()
        if pane.returncode != 0 or len(fields) != 2 or fields[0] != "1" or not fields[1].isdigit():
            die("exit-status-unavailable", 1)
        verify_live_identity(manifest, ns.socket)
        revalidate_socket_ancestry(binding)
        publish_exit_receipt(manifest_path, exit_receipt_path(ns.seat_id), ns.seat_id,
                             manifest["generation"], int(fields[1]))
        binding.close()
        return
    if tmux_args[0] == "operator-stop":
        if len(tmux_args) != 1:
            die("operator-stop-args", 1)
        verify_live_identity(manifest, ns.socket)
        revalidate_socket_ancestry(binding)
        write_stop_intent(manifest_path, manifest)
        stopped = run_tmux_bound(implementer_user, binding, ["kill-server"])
        if stopped.returncode != 0:
            die("operator-stop-tmux-failed", 1)
        verify_stopped_identity(manifest)
        publish_exit_receipt(
            manifest_path, exit_receipt_path(ns.seat_id), ns.seat_id,
            manifest["generation"], terminal_reason="operator-stop",
        )
        remove_stop_intent(manifest_path)
        binding.close()
        return
    if tmux_args[0] == "receipt":
        revalidate_socket_ancestry(binding)
        proc = run_tmux_bound(
            implementer_user, binding,
            ["list-panes", "-t", "main", "-F", "#{pane_dead} #{pane_dead_status}"],
        )
        fields = proc.stdout.strip().split()
        if proc.returncode != 0 or len(fields) != 2 or fields[0] != "1" or not fields[1].isdigit():
            die("exit-status-unavailable", 1)
        verify_live_identity(manifest, ns.socket)
        revalidate_socket_ancestry(binding)
        publish_exit_receipt(
            manifest_path,
            exit_receipt_path(ns.seat_id),
            ns.seat_id,
            manifest["generation"],
            int(fields[1]),
        )
        binding.close()
        return
    revalidate_socket_ancestry(binding)
    if tmux_args[0] == "attach-session":
        exec_tmux_attach(implementer_user, binding, tmux_args)
        die("attach-exec-returned", 1)
    proc = run_tmux_bound(implementer_user, binding, tmux_args)
    if proc.returncode != 0:
        raise SystemExit(proc.returncode)
    if proc.stdout:
        print(proc.stdout, end="")
    binding.close()


if __name__ == "__main__":
    main()