from __future__ import annotations

import hashlib
import importlib.util
import os
import sys
from pathlib import Path
from types import ModuleType
from collections.abc import Callable
from typing import Protocol, cast

import pytest


class Seal(Protocol):
    identities: dict[str, dict[str, object]]
    def close(self) -> None: ...


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: object | None = None,
    ) -> int: ...


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


def _module(path: Path, name: str) -> ModuleType:
    if path.name == "task22_launcher_runtime.py":
        modules = sys.modules
        for helper in ("task22_dependency_closure", "task22_child_protocol"):
            if helper not in modules:
                loaded = _module(ROOT / f"{helper}.py", helper)
                modules[helper] = loaded
    spec = importlib.util.spec_from_file_location(name, path)
    if spec is None or spec.loader is None:
        raise AssertionError(f"cannot load {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _launcher() -> Launcher:
    support = _module(ROOT / "task22_launcher_test_support.py", "closure_support")
    loader = cast(Callable[[Path], object], getattr(support, "load_launcher"))
    if not callable(loader):
        raise AssertionError("launcher loader missing")
    return cast(Launcher, loader(ROOT / "task22-trainer-authority-removal-canonical.py"))


@pytest.mark.parametrize("family", ("yaml", "telegram", "pydantic"))
def test_bound_dependency_replacement_after_fork_cannot_execute(
    tmp_path: Path, family: str
) -> None:
    launcher = _launcher()
    package = tmp_path / family
    _ = package.mkdir()
    original = package / "__init__.py"
    replacement = package / "replacement.py"
    marker = tmp_path / "replacement-ran"
    _ = original.write_text("BOUND = 'original'\n")
    _ = replacement.write_text(f"from pathlib import Path\nPath({str(marker)!r}).write_text('bad')\n")
    output = tmp_path / "origin"
    source = (
        "import pathlib,sys\n"
        f"sys.path.insert(0,{str(tmp_path)!r})\n"
        f"import {family} as dependency\n"
        f"pathlib.Path({str(output)!r}).write_text(dependency.__file__)\n"
    ).encode()
    descriptor = launcher.sealed_memfd("dependency-replacement", source)
    try:
        status = launcher.run_snapshot_child(
            {}, descriptor, hashlib.sha256(source).hexdigest(),
            lambda: os.replace(replacement, original),
        )
    finally:
        _ = os.close(descriptor)
    assert status == 0
    assert not marker.exists()
    assert not output.read_text().startswith(str(tmp_path))


def test_unknown_third_party_import_fails_closed(tmp_path: Path) -> None:
    launcher = _launcher()
    marker = tmp_path / "unknown-ran"
    _ = (tmp_path / "unapproved_dependency.py").write_text(
        f"from pathlib import Path\nPath({str(marker)!r}).write_text('bad')\n"
    )
    source = (
        "import sys\n"
        f"sys.path.insert(0,{str(tmp_path)!r})\n"
        "import unapproved_dependency\n"
    ).encode()
    descriptor = launcher.sealed_memfd("unknown-dependency", source)
    try:
        status = launcher.run_snapshot_child(
            {}, descriptor, hashlib.sha256(source).hexdigest(), None
        )
    finally:
        _ = os.close(descriptor)
    assert status != 0
    assert not marker.exists()


def test_runtime_binds_running_executable_and_venv_boundary() -> None:
    runtime = _module(ROOT / "task22_launcher_runtime.py", "closure_runtime")
    source = (ROOT / "task22_dependency_closure.py").read_text()
    assert "/proc/self/exe" in source
    binder = cast(Callable[[Path, Path, str, dict[str, str]], Seal], getattr(runtime, "attest_runtime"))
    assert callable(binder)
    seal = binder(
        Path("/home/cube/projects/richard/hermes-agent/.venv"),
        Path("/home/cube/miniconda3"),
        "63ee365014c377a5ea4d131ab08c31539a5734114895eb6de5b556c0e8b55ec9",
        {"pydantic": "2.13.4", "PyYAML": "6.0.3", "python-telegram-bot": "22.6"},
    )
    assert {"proc_exe", "venv_executable", "pyvenv_cfg"}.issubset(seal.identities)
    _ = seal.close()


@pytest.mark.parametrize(
    ("receipt", "status"),
    (
        (b"", 0),
        (b"\x00\x00\x00\x02{}", 0),
        (b"\x00\x00\x00\x1f{\"kind\":\"exit\",\"code\":0}", 1 << 8),
        (b"\x00\x00\x00\x21{\"kind\":\"signal\",\"signal\":15}", 0),
    ),
)
def test_missing_malformed_or_mismatched_completion_is_rejected(
    receipt: bytes, status: int
) -> None:
    runtime = _module(ROOT / "task22_launcher_runtime.py", "receipt_runtime")
    validate = cast(Callable[[bytes, int], int], getattr(runtime, "validate_completion"))
    assert callable(validate)
    with pytest.raises(RuntimeError, match="completion"):
        _ = validate(receipt, status)


def test_valid_receipt_must_match_exact_exit_or_signal_status() -> None:
    runtime = _module(ROOT / "task22_launcher_runtime.py", "status_runtime")
    protocol = _module(ROOT / "task22_child_protocol.py", "status_protocol")
    validate = cast(Callable[[bytes, int], int], getattr(runtime, "validate_completion"))
    make_frame = cast(Callable[[str, int], bytes], getattr(protocol, "frame"))
    assert callable(validate) and callable(make_frame)
    with pytest.raises(RuntimeError, match="completion.*exit status mismatch"):
        _ = validate(make_frame("exit", 0), 1 << 8)
    with pytest.raises(RuntimeError, match="completion.*status mismatch"):
        _ = validate(make_frame("signal", 15), 0)


def test_unreceipted_sigkill_is_rejected_as_incomplete() -> None:
    launcher = _launcher()
    source = b"import os,signal; os.kill(os.getpid(), signal.SIGKILL)\n"
    descriptor = launcher.sealed_memfd("unreceipted-crash", source)
    def crash() -> None:
        _ = launcher.run_snapshot_child(
            {}, descriptor, hashlib.sha256(source).hexdigest(), None
        )
    try:
        with pytest.raises(RuntimeError, match="completion"):
            crash()
    finally:
        _ = os.close(descriptor)