from __future__ import annotations

import ast
import ctypes
import hashlib
import importlib.util
import os
import select
import signal
from collections.abc import Callable
from pathlib import Path
from types import ModuleType
from typing import Protocol, cast

ROOT = Path(__file__).resolve().parent


class Launcher(Protocol):
    def sealed_memfd(self, name: str, content: bytes) -> int: ...
    def run_snapshot_child(
        self, entries: dict[str, list[object]], descriptor: int, digest: str,
        after_fork: Callable[[], None] | None = None,
        lock_handoff: tuple[Path, Path, int, dict[str, object]] | None = None,
    ) -> int: ...


def _launcher() -> Launcher:
    path = ROOT / "task22_launcher_test_support.py"
    spec = importlib.util.spec_from_file_location("death_support", path)
    if spec is None or spec.loader is None:
        raise AssertionError("support loader missing")
    module: ModuleType = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    load = cast(Callable[[Path], object], getattr(module, "load_launcher"))
    return cast(Launcher, load(ROOT / "task22-trainer-authority-removal-canonical.py"))


def _prctl(option: int, value: int) -> None:
    function = ctypes.CDLL(None, use_errno=True).prctl
    function.argtypes = [ctypes.c_int, ctypes.c_ulong, ctypes.c_ulong,
                         ctypes.c_ulong, ctypes.c_ulong]
    function.restype = ctypes.c_int
    if function(option, value, 0, 0, 0) != 0:
        error = ctypes.get_errno()
        raise OSError(error, os.strerror(error))


def _children(pid: int) -> list[int]:
    content = Path(f"/proc/{pid}/task/{pid}/children").read_text(encoding="utf-8").split()
    return [int(value) for value in content]


def _pidfd_open(pid: int) -> int:
    syscall = ctypes.CDLL(None, use_errno=True).syscall
    syscall.argtypes = [ctypes.c_long, ctypes.c_int, ctypes.c_uint]
    syscall.restype = ctypes.c_long
    descriptor = cast(int, syscall(434, pid, 0))
    if descriptor < 0:
        error = ctypes.get_errno()
        raise OSError(error, os.strerror(error))
    return descriptor


def _pidfd_event(pidfd: int, expected: int = select.POLLIN) -> None:
    poller = select.poll()
    poller.register(pidfd, select.POLLIN)
    assert poller.poll(5_000) == [(pidfd, expected)]


def _reap(pid: int) -> int:
    while True:
        try:
            waited, status = os.waitpid(pid, 0)
            assert waited == pid
            return status
        except InterruptedError:
            continue


def test_sigkill_launcher_is_supervised_without_profile_mutation_or_zombie(
    tmp_path: Path,
) -> None:
    profile = tmp_path / "profile"
    registry = profile / "customers/registry.json"
    lock = profile / "data/customer/migrations/.lock"
    marker = profile / "worker-mutated"
    registry.parent.mkdir(parents=True)
    lock.parent.mkdir(parents=True)
    _ = registry.write_bytes(b'{"customers":[]}\n')
    lock.touch(mode=0o600)
    lock.chmod(0o600)
    before = {registry: registry.read_bytes(), lock: lock.read_bytes()}
    info = lock.stat()
    expected: dict[str, object] = {
        "inode": info.st_ino, "mode": 0o600, "uid": os.geteuid(),
        "gid": os.getegid(), "size": 0,
        "sha256": hashlib.sha256(b"").hexdigest(),
    }
    ready_read, ready_write = os.pipe2(os.O_CLOEXEC)
    hold_read, hold_write = os.pipe2(os.O_CLOEXEC)
    _prctl(36, 1)  # PR_SET_CHILD_SUBREAPER
    launcher_pid = os.fork()
    if launcher_pid == 0:
        os.close(ready_read)
        os.close(hold_write)
        launcher = _launcher()
        source = f"from pathlib import Path\nPath({str(marker)!r}).write_text('bad')\n".encode()
        descriptor = launcher.sealed_memfd("parent-death-worker", source)
        lock_fd = os.open(lock, os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW)

        def blocked() -> None:
            _ = os.write(ready_write, b"1")
            _ = os.read(hold_read, 1)

        _ = launcher.run_snapshot_child(
            {}, descriptor, hashlib.sha256(source).hexdigest(), blocked,
            (profile, lock, lock_fd, expected),
        )
        os._exit(99)
    os.close(ready_write)
    os.close(hold_read)
    launcher_pidfd = _pidfd_open(launcher_pid)
    adopted: list[int] = []
    try:
        assert os.read(ready_read, 1) == b"1"
        direct = _children(launcher_pid)
        assert len(direct) == 1
        supervisor = direct[0]
        nested = _children(supervisor)
        worker = nested[0] if len(nested) == 1 else supervisor
        worker_pidfd = _pidfd_open(worker)
        supervisor_pidfd = _pidfd_open(supervisor)
        os.kill(launcher_pid, signal.SIGKILL)
        _pidfd_event(launcher_pidfd)
        assert os.WIFSIGNALED(_reap(launcher_pid))
        _pidfd_event(supervisor_pidfd)
        if supervisor != worker:
            assert os.WIFEXITED(_reap(supervisor))
            _pidfd_event(worker_pidfd, select.POLLIN | select.POLLHUP)
        else:
            _pidfd_event(worker_pidfd)
            adopted.append(worker)
        assert nested == [worker], "launcher must own a supervisor which owns the worker"
        assert not Path(f"/proc/{worker}").exists()
        assert not Path(f"/proc/{supervisor}").exists()
        assert not marker.exists()
        assert {path: path.read_bytes() for path in before} == before
        print({"launcher": launcher_pid, "supervisor": supervisor, "worker": worker,
               "tree": "launcher->supervisor->worker",
               "supervisor_event": "POLLIN", "worker_event": "POLLIN|POLLHUP",
               "supervisor_proc": False, "worker_proc": False,
               "profile_unchanged": True})
    finally:
        os.close(hold_write)
        os.close(ready_read)
        os.close(launcher_pidfd)
        for descriptor in (locals().get("worker_pidfd"), locals().get("supervisor_pidfd")):
            if isinstance(descriptor, int):
                os.close(descriptor)
        for pid in adopted:
            _ = _reap(pid)
        _prctl(36, 0)


def test_lifecycle_tests_use_no_sleep_or_process_whitelist() -> None:
    tree = ast.parse(Path(__file__).read_text(encoding="utf-8"))
    attributes = [node.attr for node in ast.walk(tree) if isinstance(node, ast.Attribute)]
    assert "sleep" not in attributes
    assert "process_iter" not in attributes