"""Cross-turn broker for ManagedLockSet issue/repository FLOCKs."""
from __future__ import annotations

import argparse
import errno
import json
import os
import secrets
import signal
import socket
import stat
import sys
import time
from pathlib import Path

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from harness.lib.managed_lock import LockTimeoutError, ManagedLockSet, probe_free

ROLES = {"gitea-impl-agent", "gitea-spec-agent", "gitea-review-agent"}
MAX_TTL = 3600


def _abs(value: str, name: str) -> Path:
    p = Path(value)
    if not p.is_absolute() or "\x00" in value:
        raise ValueError(f"{name} must be absolute")
    return p


def _safe_dir(path: Path, name: str) -> None:
    try:
        fd = _open_directory(path)
    except OSError as exc:
        raise ValueError(f"unsafe {name} parent") from exc
    try:
        st = os.fstat(fd)
        if (not stat.S_ISDIR(st.st_mode) or st.st_uid != os.getuid()
                or stat.S_IMODE(st.st_mode) != 0o700):
            raise ValueError(f"unsafe {name} parent")
    finally:
        os.close(fd)


def _open_directory(path: Path) -> int:
    parts = path.parts
    fd = os.open(parts[0], os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC)
    try:
        for part in parts[1:]:
            nxt = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC,
                          dir_fd=fd)
            os.close(fd)
            fd = nxt
        return fd
    except BaseException:
        os.close(fd)
        raise


def _open_parent(path: Path) -> tuple[int, str]:
    return _open_directory(path.parent), path.name


def _validate_handle_fd(fd: int, path: Path) -> os.stat_result:
    fd_st = os.fstat(fd)
    if (not stat.S_ISREG(fd_st.st_mode)
            or fd_st.st_uid != os.getuid()
            or stat.S_IMODE(fd_st.st_mode) != 0o600
            or fd_st.st_nlink != 1):
        raise ValueError("unsafe handle file")
    return fd_st


def _open_handle(path: Path, flags: int) -> int:
    parent_fd, name = _open_parent(path)
    try:
        fd = os.open(name, flags | os.O_NOFOLLOW | os.O_CLOEXEC, dir_fd=parent_fd)
    except OSError as exc:
        if exc.errno in (errno.ELOOP, errno.ENOTDIR):
            raise ValueError("unsafe handle file") from exc
        raise
    try:
        _validate_handle_fd(fd, path)
        return fd
    except BaseException:
        os.close(fd)
        raise
    finally:
        os.close(parent_fd)


def _safe_existing_regular(path: Path) -> None:
    fd = _open_handle(path, os.O_RDONLY)
    os.close(fd)


def _write_all(fd: int, data: bytes) -> None:
    written = 0
    while written < len(data):
        written += os.write(fd, data[written:])


def _write_handle(path: Path, state: dict, *, exclusive: bool = False, on_created=None) -> None:
    flags = os.O_WRONLY | os.O_CREAT | (os.O_EXCL if exclusive else 0)
    if not exclusive:
        fd = _open_handle(path, flags)
        parent_fd = None
    else:
        parent_fd, name = _open_parent(path)
        fd = None
        try:
            fd = os.open(name, flags | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600, dir_fd=parent_fd)
            created = os.fstat(fd)
            if on_created is not None:
                on_created((created.st_dev, created.st_ino))
        except BaseException:
            if fd is not None:
                os.close(fd)
            raise
        finally:
            os.close(parent_fd)
    try:
        if exclusive:
            os.fchmod(fd, 0o600)
        _validate_handle_fd(fd, path)
        os.lseek(fd, 0, os.SEEK_SET)
        os.ftruncate(fd, 0)
        _write_all(fd, (json.dumps(state, sort_keys=True, separators=(",", ":")) + "\n").encode())
        os.fsync(fd)
    finally:
        os.close(fd)


def _update_handle(path: Path, state: dict, *, expected_identity: tuple[int, int] | None = None) -> None:
    fd = _open_handle(path, os.O_WRONLY)
    try:
        st = os.fstat(fd)
        if expected_identity is not None and (st.st_dev, st.st_ino) != expected_identity:
            raise ValueError("handle identity changed")
        os.lseek(fd, 0, os.SEEK_SET)
        os.ftruncate(fd, 0)
        _write_all(fd, (json.dumps(state, sort_keys=True, separators=(",", ":")) + "\n").encode())
        os.fsync(fd)
    finally:
        os.close(fd)


def _read_handle(path: Path) -> dict:
    fd = _open_handle(path, os.O_RDONLY)
    try:
        chunks = []
        while True:
            chunk = os.read(fd, 65536)
            if not chunk:
                break
            chunks.append(chunk)
        data = json.loads(b"".join(chunks).decode())
    finally:
        os.close(fd)
    if not isinstance(data, dict) or not isinstance(data.get("token"), str) or not isinstance(data.get("socket"), str):
        raise ValueError("invalid handle")
    return data


def read_handle_metadata(path: os.PathLike[str] | str) -> dict:
    """Return non-secret, authenticated-handle identity metadata only."""
    data = _read_handle(_abs(str(path), "handle"))
    return {key: value for key, value in data.items() if key != "token"}


class Broker:
    """Broker locks are trusted by OS UID; random token authenticates the holder.

    Same-UID role strings are metadata, not authorization: the same OS UID is the
    security boundary and an unpredictable token rejects non-holder requests.
    """
    def __init__(self, args: argparse.Namespace):
        if args.role not in ROLES:
            raise ValueError("invalid role")
        if not isinstance(args.run_id, str) or not args.run_id or "\x00" in args.run_id or "/" in args.run_id or "\\" in args.run_id:
            raise ValueError("invalid run_id")
        if not isinstance(args.issue, int) or args.issue <= 0:
            raise ValueError("issue must be positive")
        if not 1 <= args.ttl_seconds <= MAX_TTL:
            raise ValueError("ttl_seconds out of bounds")
        self.root = _abs(args.root, "root")
        self.handle = _abs(args.handle, "handle")
        self.sock_path = _abs(args.socket, "socket")
        locks = self.root / "repositories" / args.owner / args.repository / "locks"
        _safe_dir(locks, "repository locks")
        _safe_dir(self.handle.parent, "handle")
        _safe_dir(self.sock_path.parent, "socket")
        self.token = secrets.token_hex(32)
        self.deadline = time.monotonic() + args.ttl_seconds
        self.set = ManagedLockSet.for_issue(self.root, args.owner, args.repository, args.issue, timeout=args.acquire_timeout)
        self.meta = {"root": str(self.root), "owner": args.owner, "repository": args.repository,
                     "issue": args.issue, "role": args.role, "run_id": args.run_id,
                     "pid": os.getpid(), "deadline": self.deadline}
        self.state = {**self.meta, "socket": str(self.sock_path), "status": "held", "token": self.token}
        self.server: socket.socket | None = None
        self.socket_identity: tuple[int, int] | None = None
        self.socket_parent_fd: int | None = None
        self.stopping = False
        self.handle_created = False
        self.handle_identity: tuple[int, int] | None = None
        self.started = False
        self.cleanup_result: dict | None = None

    def _record_handle_identity(self, identity: tuple[int, int]) -> None:
        self.handle_identity = identity
        self.handle_created = True

    def persist(self, status: str, *, retain_token: bool = True) -> None:
        self.state["status"] = status
        self.state["deadline"] = self.deadline
        persisted = dict(self.state)
        if not retain_token:
            persisted.pop("token", None)
        _update_handle(self.handle, persisted, expected_identity=self.handle_identity)

    def cleanup(self, status: str) -> bool:
        # Keep a previous partial result: a second cleanup must not erase failures.
        previous = self.cleanup_result or {}
        errors = list(previous.get("errors", []))
        for name, release in (("repository-release", self.set.repository.release),
                              ("issue-release", self.set.release_issue)):
            try:
                if name == "issue-release" or self.set.repository.held:
                    release()
            except Exception:
                errors.append(name)
        released = not self.set.issue.held and not self.set.repository.held
        if self.server:
            server, self.server = self.server, None
            try:
                server.close()
            except OSError:
                errors.append("server-close")
        socket_removed = previous.get("socket_removed", self.socket_identity is None)
        if self.socket_identity is not None and self.socket_parent_fd is not None:
            try:
                st = os.stat(self.sock_path.name, dir_fd=self.socket_parent_fd, follow_symlinks=False)
                if (not stat.S_ISSOCK(st.st_mode) or st.st_uid != os.getuid()
                        or (st.st_dev, st.st_ino) != self.socket_identity):
                    errors.append("socket-identity-changed")
                else:
                    os.unlink(self.sock_path.name, dir_fd=self.socket_parent_fd)
                    socket_removed = True
                    self.socket_identity = None
            except FileNotFoundError:
                socket_removed = True
                self.socket_identity = None
            except OSError:
                errors.append("socket-unlink")
        if self.socket_parent_fd is not None:
            fd, self.socket_parent_fd = self.socket_parent_fd, None
            try:
                os.close(fd)
            except OSError:
                errors.append("socket-parent-close")
        self.cleanup_result = {"lease_released": released, "socket_removed": bool(socket_removed),
                               "handle_saved": False, "errors": sorted(set(errors))}
        if self.handle_created:
            try:
                # Do not claim another object's handle during failed startup.
                if not self.started and self.handle_identity is not None:
                    fd = _open_handle(self.handle, os.O_RDONLY)
                    try:
                        st = os.fstat(fd)
                        if (st.st_dev, st.st_ino) != self.handle_identity:
                            raise ValueError("handle identity changed")
                    finally:
                        os.close(fd)
                self.cleanup_result["handle_saved"] = True
                self.state["cleanup"] = self.cleanup_result
                terminal_status = status if released and socket_removed and not errors else "cleanup-incomplete"
                self.persist(terminal_status, retain_token=not released)
            except Exception:
                self.cleanup_result["handle_saved"] = False
                self.cleanup_result["errors"] = sorted(set(self.cleanup_result["errors"] + ["handle-save"]))
        else:
            self.cleanup_result["handle_saved"] = True  # No owned handle to save.
        if self.cleanup_result["errors"]:
            print(json.dumps({"event": "cleanup-incomplete", "cleanup": self.cleanup_result},
                             separators=(",", ":")), file=sys.stderr, flush=True)
        return released and bool(socket_removed) and not self.cleanup_result["errors"]

    def authenticate(self, req: dict) -> bool:
        supplied = req.get("token")
        return isinstance(supplied, str) and secrets.compare_digest(supplied, self.token)

    def handle_request(self, req: dict) -> dict:
        if not self.authenticate(req):
            return {"error": "unauthorized"}
        if time.monotonic() >= self.deadline:
            self.cleanup("expired")
            self.stopping = True
            return {"error": "expired"}
        action = req.get("action")
        if action == "status":
            return {"status": self.state["status"], "deadline": self.deadline,
                    "repository_held": self.set.repository.held, "root": self.state["root"],
                    "owner": self.state["owner"], "repository": self.state["repository"],
                    "issue": self.state["issue"], "role": self.state["role"],
                    "pid": self.state["pid"]}
        if action == "renew":
            ttl = req.get("ttl_seconds")
            if isinstance(ttl, bool) or not isinstance(ttl, (int, float)) or not 1 <= ttl <= MAX_TTL:
                return {"error": "invalid ttl_seconds"}
            self.deadline = time.monotonic() + ttl
            self.persist("repository-held" if self.set.repository.held else "held")
            return {"status": self.state["status"], "deadline": self.deadline}
        if action == "acquire_repository":
            try:
                self.set.repository.acquire()
            except LockTimeoutError:
                return {"error": "repository-busy"}
            self.persist("repository-held")
            return {"status": "repository-held"}
        if action == "release_repository":
            if self.set.repository.held:
                self.set.repository.release()
            self.persist("held")
            return {"status": "held"}
        if action == "release":
            complete = self.cleanup("released")
            self.stopping = True
            if not complete:
                return {"error": "cleanup-incomplete", "cleanup": self.cleanup_result}
            return {"status": "released"}
        return {"error": "invalid action"}

    def run(self) -> None:
        try:
            self.set.issue.acquire()
            signal.signal(signal.SIGTERM, lambda *_: setattr(self, "stopping", True))
            signal.signal(signal.SIGINT, lambda *_: setattr(self, "stopping", True))
            self.set.issue.write_metadata(self.meta)
            try:
                _write_handle(self.handle, self.state, exclusive=True,
                              on_created=self._record_handle_identity)
            except BaseException:
                raise
            st = self.handle.lstat()
            if (st.st_dev, st.st_ino) != self.handle_identity:
                raise ValueError("created handle identity changed")
            self.handle_created = True
            self.socket_parent_fd, socket_name = _open_parent(self.sock_path)
            try:
                os.stat(socket_name, dir_fd=self.socket_parent_fd, follow_symlinks=False)
                raise ValueError("socket path already exists")
            except FileNotFoundError:
                pass
            self.server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            self.server.bind(f"/proc/self/fd/{self.socket_parent_fd}/{socket_name}")
            # Socket FD metadata describes sockfs, not the bound filesystem node.
            sock_st = os.stat(socket_name, dir_fd=self.socket_parent_fd, follow_symlinks=False)
            if not stat.S_ISSOCK(sock_st.st_mode) or sock_st.st_uid != os.getuid():
                raise ValueError("unsafe bound socket")
            self.socket_identity = (sock_st.st_dev, sock_st.st_ino)
            os.chmod(socket_name, 0o600, dir_fd=self.socket_parent_fd, follow_symlinks=False)
            verified = os.stat(socket_name, dir_fd=self.socket_parent_fd, follow_symlinks=False)
            if (not stat.S_ISSOCK(verified.st_mode) or verified.st_uid != os.getuid()
                    or (verified.st_dev, verified.st_ino) != self.socket_identity
                    or stat.S_IMODE(verified.st_mode) != 0o600):
                raise ValueError("bound socket verification failed")
            self.server.listen(8); self.server.settimeout(0.2)
            self.started = True
            print(json.dumps({"status": "held", "handle": str(self.handle), "socket": str(self.sock_path)}, separators=(",", ":")), flush=True)
            while not self.stopping:
                if time.monotonic() >= self.deadline:
                    self.stopping = True
                    self.cleanup("expired"); break
                try:
                    conn, _ = self.server.accept()
                except socket.timeout:
                    continue
                with conn:
                    conn.settimeout(1)
                    line = conn.makefile("rb").readline(65536)
                    try:
                        req = json.loads(line.decode())
                        response = self.handle_request(req) if isinstance(req, dict) else {"error": "invalid request"}
                    except Exception:
                        response = {"error": "invalid request"}
                    conn.sendall((json.dumps(response, separators=(",", ":")) + "\n").encode())
        finally:
            if (self.set.issue.held or self.set.repository.held
                    or (self.cleanup_result is None and (self.handle_created or self.server))):
                self.cleanup("released" if not self.stopping else "expired")
        if self.cleanup_result and self.cleanup_result["errors"]:
            raise OSError("broker cleanup incomplete")


def serve(args: argparse.Namespace) -> int:
    broker = Broker(args)
    broker.run()
    return 0


def request_handle(handle_path: str | Path, action: str, ttl: float | None = None) -> dict:
    """Send an authenticated request using a validated handle, without printing."""
    handle = _abs(str(handle_path), "handle")
    data = _read_handle(handle)
    sock_path = _abs(data["socket"], "socket")
    req = {"action": action, "token": data["token"]}
    if ttl is not None:
        req["ttl_seconds"] = ttl
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
        sock.settimeout(3); sock.connect(str(sock_path)); sock.sendall((json.dumps(req) + "\n").encode())
        with sock.makefile("rb") as response_file:
            response = json.loads(response_file.readline())
    if not isinstance(response, dict):
        raise ValueError("invalid broker response")
    return response


def request(args: argparse.Namespace) -> int:
    response = request_handle(args.handle, args.action, args.ttl_seconds)
    print(json.dumps(response, separators=(",", ":")))
    return 0 if "error" not in response else 1


def probe(args: argparse.Namespace) -> int:
    root = _abs(args.root, "root")
    issue_path, repository_path = ManagedLockSet.paths_for_issue(
        root, args.owner, args.repository, args.issue)
    _safe_dir(issue_path.parent, "repository locks")
    issue_free = probe_free(issue_path)
    repository_free = probe_free(repository_path)
    if args.resource == "issue":
        free = issue_free
    elif args.resource == "repository":
        free = repository_free
    else:
        free = issue_free and repository_free
    print(json.dumps({"free": free, "issue_free": issue_free,
                      "repository_free": repository_free}, separators=(",", ":")))
    return 0 if free else 1


def parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="command", required=True)
    s = sub.add_parser("serve")
    s.add_argument("--root", required=True); s.add_argument("--owner", required=True); s.add_argument("--repository", required=True); s.add_argument("--issue", required=True, type=int); s.add_argument("--role", required=True); s.add_argument("--run-id", required=True); s.add_argument("--handle", required=True); s.add_argument("--socket", required=True); s.add_argument("--ttl-seconds", required=True, type=float); s.add_argument("--acquire-timeout", type=float, default=0.0)
    r = sub.add_parser("request")
    r.add_argument("--handle", required=True); r.add_argument("--action", required=True, choices=["status", "renew", "acquire_repository", "release_repository", "release"]); r.add_argument("--ttl-seconds", type=float)
    check = sub.add_parser("probe")
    check.add_argument("--root", required=True); check.add_argument("--owner", required=True); check.add_argument("--repository", required=True); check.add_argument("--issue", required=True, type=int)
    check.add_argument("--resource", choices=["issue", "repository", "all"], default="all")
    return p


def main() -> int:
    try:
        args = parser().parse_args()
        if args.command == "serve":
            return serve(args)
        if args.command == "request":
            return request(args)
        return probe(args)
    except (ValueError, OSError, LockTimeoutError, json.JSONDecodeError, KeyError, TypeError) as exc:
        print(f"error: {type(exc).__name__}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())