'''
Unit-ish tests for specific IPC transport protocol backends.

'''
from __future__ import annotations
import os
from pathlib import Path
import socket
import stat
import struct
import sys
import tempfile
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import trio
from trio.testing import (
    MockClock,
    wait_all_tasks_blocked,
)
import tractor
from tractor import Actor
from tractor.discovery import _addr
from tractor.ipc._transport import MsgpackTransport
from tractor.runtime import _state



def test_cancelled_transport_send_completes_frame():
    '''
    Finish an in-flight frame before delivering sender cancellation.

    A cancelled `send_all()` may leave an arbitrary frame prefix on the
    wire. Closing the actor-wide stream avoids decoder corruption but
    also destroys unrelated contexts using that channel. On its first
    call, the fake stream publishes two header bytes and blocks until
    the parent test releases it. This lets the parent cancel the sender
    while frame publication is suspended. The sender must remain inside
    `send_first()` until the complete frame is written, then observe
    pending cancellation; a second sender proves sibling contexts can
    safely reuse the still frame-aligned stream.

    '''
    class PartialSendStream:
        def __init__(self) -> None:
            self.send_all_entered = trio.Event()
            self.send_all_release = trio.Event()
            self.closed = False
            self.wire = bytearray()

        async def send_all(
            self,
            data: bytes,
        ) -> None:
            assert data
            if not self.wire:
                self.wire.extend(data[:2])
                self.send_all_entered.set()
                await self.send_all_release.wait()
                self.wire.extend(data[2:])
            else:
                self.wire.extend(data)

        async def aclose(self) -> None:
            self.closed = True

    def count_frames(wire: bytearray) -> int:
        offset: int = 0
        count: int = 0
        while offset < len(wire):
            header_end: int = offset + 4
            assert header_end <= len(wire)
            size, = struct.unpack('<I', wire[offset:header_end])
            offset = header_end + size
            assert offset <= len(wire)
            count += 1

        assert offset == len(wire)
        return count

    async def main() -> None:
        stream = PartialSendStream()
        transport = object.__new__(MsgpackTransport)
        transport.stream = stream
        transport._send_lock = trio.StrictFIFOLock()
        sender_done = trio.Event()
        sender_scopes: list[trio.CancelScope] = []
        cancelled_caught: bool = False

        first_msg = tractor.msg.Start(
            ns=__name__,
            func='add_one',
            kwargs={'n': 1},
            uid=('root', 'test'),
            cid='partial-send',
        )
        second_msg = tractor.msg.Start(
            ns=__name__,
            func='add_one',
            kwargs={'n': 2},
            uid=('root', 'test'),
            cid='second-send',
        )

        async def send_first() -> None:
            nonlocal cancelled_caught
            with trio.CancelScope() as cs:
                sender_scopes.append(cs)
                await transport.send(first_msg)

            cancelled_caught = cs.cancelled_caught
            sender_done.set()

        async with trio.open_nursery() as tn:
            tn.start_soon(
                send_first,
            )
            await stream.send_all_entered.wait()
            sender_scopes[0].cancel()
            await wait_all_tasks_blocked()

            assert not stream.closed
            # Cancellation is pending, but complete-frame shielding
            # keeps `send_first()` suspended in `.send_all()`.
            assert not sender_done.is_set()

            # Let the underlying frame write finish after the parent
            # has requested sender cancellation.
            stream.send_all_release.set()
            await sender_done.wait()

            assert cancelled_caught
            assert not stream.closed
            # The initial two-byte prefix was completed into one valid
            # frame before cancellation reached `send_first()`.
            assert count_frames(stream.wire) == 1

            await transport.send(second_msg)
            # A sibling sender can append and decode another frame only
            # because the first cancellation preserved stream alignment.
            assert count_frames(stream.wire) == 2

            tn.cancel_scope.cancel()

    trio.run(main)


def test_transport_send_deadline_closes_partial_frame():
    '''
    Destroy a stalled partial frame before another sender can append.

    Ordinary cancellation cannot interrupt complete-frame publication.
    Bounded actor/context cancellation instead passes its absolute
    deadline into this operation. The fake stream writes a partial
    header and stalls; when the send's own deadline fires, the transport
    must close the stream before releasing its shared send lock. This
    prevents the next sender from appending bytes which a decoder would
    treat as the remainder of the corrupt first frame.

    '''
    class StalledStream:
        def __init__(self) -> None:
            self.closed = False
            self.wire = bytearray()

        async def send_all(
            self,
            data: bytes,
        ) -> None:
            self.wire.extend(data[:2])
            await trio.sleep_forever()

        async def aclose(self) -> None:
            self.closed = True

    async def main() -> None:
        stream = StalledStream()
        transport = object.__new__(MsgpackTransport)
        transport.stream = stream
        transport._send_lock = trio.StrictFIFOLock()
        msg = tractor.msg.Start(
            ns=__name__,
            func='add_one',
            kwargs={'n': 1},
            uid=('root', 'test'),
            cid='deadline-send',
        )

        with pytest.raises(
            tractor.TransportClosed,
            match='frame publication exceeded',
        ):
            await transport.send(
                msg,
                send_deadline=1,
            )

        assert stream.closed  # partial-frame timeout destroys stream
        assert len(stream.wire) == 2  # only a header fragment was sent
        assert not transport._send_lock.locked()  # cleanup released lock

    trio.run(
        main,
        clock=MockClock(autojump_threshold=0),
    )


def test_cancelled_transport_send_preserves_cancellation():
    '''
    Prefer sender cancellation when teardown closes the stream.

    `MsgpackTransport.send()` shields frame publication at
    `tractor.ipc._transport:MsgpackTransport.send`. Before this fix, an
    outer `move_on_after()`/cancel scope could cancel `Channel.send()`
    while actor teardown closed the shared stream. The resulting
    `ClosedResourceError` escaped from the transport handler instead of
    its `checkpoint_if_cancelled()` redelivering pending cancellation.

    The fake stream blocks inside the shield until the test cancels the
    sender. Parent-controlled release then simulates actor teardown
    closing the socket and raises the `ClosedResourceError` observed on
    macOS UDS. `CancelScope.cancelled_caught` proves the handler's
    checkpoint preserved cancellation as the primary outcome instead of
    leaking that secondary close error.

    '''
    class ClosingStream:
        def __init__(self) -> None:
            self.send_all_entered = trio.Event()
            self.send_all_release = trio.Event()

        async def send_all(
            self,
            data: bytes,
        ) -> None:
            assert data
            self.send_all_entered.set()
            await self.send_all_release.wait()
            # Model actor teardown closing the shared transport while
            # this sender is still inside the complete-frame shield.
            raise trio.ClosedResourceError(
                'this socket was already closed'
            )

    async def main() -> None:
        stream = ClosingStream()
        transport = object.__new__(MsgpackTransport)
        transport.stream = stream
        transport._send_lock = trio.StrictFIFOLock()
        sender_done = trio.Event()
        sender_scopes: list[trio.CancelScope] = []
        cancelled_caught: bool = False

        msg = tractor.msg.Start(
            ns=__name__,
            func='add_one',
            kwargs={'n': 1},
            uid=('root', 'test'),
            cid='close-during-cancelled-send',
        )

        async def send() -> None:
            nonlocal cancelled_caught
            with trio.CancelScope() as cs:
                sender_scopes.append(cs)
                await transport.send(msg)

            cancelled_caught = cs.cancelled_caught
            sender_done.set()

        async with trio.open_nursery() as tn:
            tn.start_soon(send)
            await stream.send_all_entered.wait()
            sender_scopes[0].cancel()
            await wait_all_tasks_blocked()

            assert not sender_done.is_set()
            stream.send_all_release.set()
            await sender_done.wait()

        assert cancelled_caught

    trio.run(main)


@pytest.fixture
def bindspace_dir_str() -> str:

    from tractor.runtime._state import get_rt_dir
    rt_dir: Path = get_rt_dir()
    bs_dir: Path = rt_dir / 'doggy'
    bs_dir_str: str = str(bs_dir)
    assert not bs_dir.is_dir()

    yield bs_dir_str

    # delete it on suite teardown.
    # ?TODO? should we support this internally
    # or is leaking it ok?
    if bs_dir.is_dir():
        bs_dir.rmdir()


def test_macos_rt_dir_fits_uds_path_limit(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    '''
    Keep the default Darwin UDS bindpath below its 104-byte limit.

    `platformdirs` normally places the runtime directory below the
    long `~/Library/Caches/TemporaryItems` path. Pytest also assigns
    a deeply nested temporary home, so appending a registry socket
    name made every macOS UDS listener fail with `AF_UNIX path too
    long`. This test simulates Darwin and an intentionally long
    platformdirs result, then proves `get_rt_dir()` uses the short
    system temporary directory and leaves room for the socket name.

    '''
    long_rt_dir: Path = tmp_path / ('long' * 30)
    monkeypatch.setattr(sys, 'platform', 'darwin')
    monkeypatch.setattr(
        'platformdirs.user_runtime_dir',
        lambda appname: str(long_rt_dir / appname),
    )
    monkeypatch.setattr(_state, '_DARWIN_TMPDIR', tmp_path)
    rt_dir: Path = _state.get_rt_dir()
    sockpath: Path = (
        Path('/tmp')
        / f'tractor-{os.getuid()}'
        / 'registry@1616.sock'
    )

    assert rt_dir == tmp_path / f'tractor-{os.getuid()}'
    assert len(os.fsencode(sockpath)) < 104
    assert stat.S_IMODE(rt_dir.stat().st_mode) == 0o700


def test_macos_rt_dir_rejects_symlink(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    '''
    Reject a pre-created symlink at the Darwin runtime path.

    Darwin uses the predictable `/tmp/tractor-<uid>` path to stay
    below its `AF_UNIX` limit. A hostile local user could otherwise
    point that path at a victim-owned directory and make
    `get_rt_dir()` chmod or place sockets in the symlink target. The
    test replaces `/tmp` with a controlled directory, installs the
    malicious link, and proves non-following validation rejects it.

    '''
    runtime_link: Path = tmp_path / f'tractor-{os.getuid()}'
    target_dir: Path = tmp_path / 'target'
    target_dir.mkdir(mode=0o755)
    runtime_link.symlink_to(target_dir, target_is_directory=True)
    monkeypatch.setattr(sys, 'platform', 'darwin')
    monkeypatch.setattr(_state, '_DARWIN_TMPDIR', tmp_path)

    with pytest.raises(PermissionError, match='Unsafe Darwin'):
        _state.get_rt_dir()

    assert stat.S_IMODE(target_dir.stat().st_mode) == 0o755


def test_reaper_uses_default_uds_bindspace(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    '''
    Sweep the same platform-specific bindspace used by UDS actors.

    The reaper previously consulted only `XDG_RUNTIME_DIR`, missing
    Darwin sockets after the runtime moved to `/tmp/tractor-<uid>`.
    This test replaces `UDSAddress.def_bindspace` and proves the test
    harness resolves that shared transport default directly.

    '''
    from tractor._testing import _reap
    from tractor.ipc._uds import UDSAddress

    monkeypatch.setattr(
        UDSAddress,
        'def_bindspace',
        tmp_path,
    )

    assert _reap.get_uds_dir() == str(tmp_path)


def test_automatic_reaper_preserves_registry_sentinel(
    monkeypatch: pytest.MonkeyPatch,
):
    '''
    Reserve unconditional registry cleanup for the explicit CLI.

    The `registry@1616.sock` suffix does not encode its binder PID, so
    automatic pytest cleanup cannot distinguish a leak from another
    live registrar. This test creates registry and actor sockets,
    proves the default sweep selects only the dead actor, then proves
    explicit sentinel inclusion retains the CLI's documented behavior.

    '''
    from tractor._testing import _reap

    with tempfile.TemporaryDirectory(
        prefix='tractor-reap-',
        dir='/tmp',
    ) as tmpdir:
        bindspace: Path = Path(tmpdir)
        registry_path: Path = bindspace / 'registry@1616.sock'
        actor_path: Path = bindspace / 'worker@1234.sock'
        socks: list[socket.socket] = []
        for path in (registry_path, actor_path):
            sock = socket.socket(socket.AF_UNIX)
            sock.bind(str(path))
            socks.append(sock)

        monkeypatch.setattr(_reap, '_is_alive', lambda pid: False)
        try:
            assert _reap.find_orphaned_uds(
                uds_dir=str(bindspace),
            ) == [str(actor_path)]
            assert set(
                _reap.find_orphaned_uds(
                    uds_dir=str(bindspace),
                    include_registry_sentinel=True,
                )
            ) == {
                str(registry_path),
                str(actor_path),
            }
        finally:
            for sock in socks:
                sock.close()


def test_rt_dir_rejects_non_directory(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    '''
    Preserve the non-Darwin runtime-directory type contract.

    Replacing `Path.is_dir()` with unguarded `lstat()` briefly made
    existing files look like valid runtime directories on Linux.
    This test points `platformdirs` at a regular file and proves
    `get_rt_dir()` rejects it during initialization.

    '''
    rt_file: Path = tmp_path / 'runtime-file'
    rt_file.touch()
    monkeypatch.setattr(sys, 'platform', 'linux')
    monkeypatch.setattr(
        'platformdirs.user_runtime_dir',
        lambda appname: str(rt_file),
    )

    with pytest.raises(
        PermissionError,
        match='Unsafe POSIX',
    ):
        _state.get_rt_dir()

    new_rt_dir: Path = tmp_path / 'new-runtime-dir'
    monkeypatch.setattr(
        'platformdirs.user_runtime_dir',
        lambda appname: str(new_rt_dir),
    )
    assert _state.get_rt_dir() == new_rt_dir
    assert stat.S_IMODE(new_rt_dir.stat().st_mode) == 0o700


def test_linux_rt_dir_secures_existing_path(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    '''
    Enforce owner-only access on an existing Linux runtime directory.

    Linux previously accepted any existing directory returned by
    `platformdirs`, without checking ownership or correcting a
    traversable mode. This test creates an owner-controlled `0o755`
    directory and proves `get_rt_dir()` normalizes the managed
    bindspace to `0o700` before returning it.

    '''
    rt_dir: Path = tmp_path / 'tractor'
    rt_dir.mkdir(mode=0o755)
    monkeypatch.setattr(sys, 'platform', 'linux')
    monkeypatch.setattr(
        'platformdirs.user_runtime_dir',
        lambda appname: str(rt_dir),
    )

    assert _state.get_rt_dir() == rt_dir
    assert stat.S_IMODE(rt_dir.stat().st_mode) == 0o700


def test_linux_rt_dir_rejects_foreign_owner(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    '''
    Reject an existing Linux runtime directory owned by another UID.

    A pre-created bindspace must never be made private with `chmod`
    until ownership is verified. This test makes the current process
    appear to have a different UID and proves `get_rt_dir()` rejects
    the directory without changing its original mode.

    '''
    rt_dir: Path = tmp_path / 'tractor'
    rt_dir.mkdir(mode=0o755)
    original_mode: int = stat.S_IMODE(rt_dir.stat().st_mode)
    monkeypatch.setattr(sys, 'platform', 'linux')
    monkeypatch.setattr(
        'platformdirs.user_runtime_dir',
        lambda appname: str(rt_dir),
    )
    monkeypatch.setattr(
        os,
        'getuid',
        lambda: rt_dir.stat().st_uid + 1,
    )

    with pytest.raises(
        PermissionError,
        match='Unsafe POSIX',
    ):
        _state.get_rt_dir()

    assert stat.S_IMODE(rt_dir.stat().st_mode) == original_mode


def test_macos_rt_dir_rejects_intermediate_symlink(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    '''
    Reject symlinks in nested Darwin runtime subdirectories.

    The earlier final-component check allowed `link/child` to follow
    an intermediate symlink and create `child` outside the secured
    runtime root. This test installs that link and proves traversal
    stops before anything is created in its target.

    '''
    rt_root: Path = tmp_path / f'tractor-{os.getuid()}'
    target_dir: Path = tmp_path / 'target'
    rt_root.mkdir(mode=0o700)
    target_dir.mkdir()
    (rt_root / 'link').symlink_to(
        target_dir,
        target_is_directory=True,
    )
    monkeypatch.setattr(sys, 'platform', 'darwin')
    monkeypatch.setattr(_state, '_DARWIN_TMPDIR', tmp_path)

    with pytest.raises(PermissionError, match='Unsafe Darwin'):
        _state.get_rt_dir(subdir='link/child')

    assert not (target_dir / 'child').exists()


@pytest.mark.parametrize(
    ('platform_name', 'path_limit'),
    [
        ('darwin', 104),
        ('linux', 108),
    ],
)
def test_uds_sockname_compaction(
    monkeypatch: pytest.MonkeyPatch,
    platform_name: str,
    path_limit: int,
):
    '''
    Keep generated actor sockets safe and below Darwin's byte limit.

    Actor names are unrestricted identity strings. A long, multibyte,
    or path-like name previously produced overlong or escaping socket
    paths. These cases prove `UDSAddress.get_sockname()` preserves a
    short legacy name, deterministically compacts unsafe names, keeps
    the reaper's `@pid.sock` suffix, and stays within Darwin's byte
    limit.

    '''
    from tractor.ipc._uds import UDSAddress

    bindspace: Path = Path('/tmp/tractor-501')
    pid: int = 12345
    from tractor.ipc import _uds

    monkeypatch.setattr(sys, 'platform', platform_name)
    monkeypatch.setattr(_uds, '_SUN_PATH_LIMIT', path_limit)

    short: Path = UDSAddress.get_sockname(
        name='worker',
        pid=pid,
        bindspace=bindspace,
    )
    long_name: str = 'actor-' + ('\u00e9' * 100)
    compact: Path = UDSAddress.get_sockname(
        name=long_name,
        pid=pid,
        bindspace=bindspace,
    )
    unsafe: Path = UDSAddress.get_sockname(
        name='../worker',
        pid=pid,
        bindspace=bindspace,
    )

    assert short == Path(f'worker@{pid}.sock')
    assert compact == UDSAddress.get_sockname(
        name=long_name,
        pid=pid,
        bindspace=bindspace,
    )
    assert compact.name.endswith(f'@{pid}.sock')
    assert unsafe.parent == Path('.')
    assert '..' not in unsafe.name
    assert len(os.fsencode(bindspace / compact)) < path_limit

    with pytest.raises(ValueError) as exc_info:
        UDSAddress.get_sockname(
            name=long_name,
            pid=pid,
            bindspace=Path('/tmp') / ('x' * 90),
        )

    errmsg: str = str(exc_info.value)
    assert 'leaves no room' in errmsg
    assert 'name was unsafe: False' in errmsg
    assert 'name was over budget: True' in errmsg
    assert f'AF_UNIX path limit: {path_limit}' in errmsg


def test_uds_reaper_ignores_unreconstructable_path(
    monkeypatch: pytest.MonkeyPatch,
):
    '''
    Keep post-kill UDS cleanup best-effort on path overflow.

    `unlink_uds_bind_addrs()` reconstructs a self-assigned socket from
    the dead actor's name and PID. An over-budget bindspace makes that
    naming helper raise before `os.unlink()`; propagating the error
    would replace the original supervision outcome after the child was
    already killed. This test forces overflow and proves cleanup skips
    reconstruction without attempting an unlink or raising.

    '''
    from tractor.ipc import _uds
    from tractor.spawn import _reap

    long_bindspace: Path = Path('/tmp') / ('x' * 120)
    proc = SimpleNamespace(pid=12345)
    subactor = SimpleNamespace(
        aid=SimpleNamespace(name='worker'),
    )
    unlink = Mock()
    monkeypatch.setattr(
        _uds.UDSAddress,
        'def_bindspace',
        long_bindspace,
    )
    monkeypatch.setattr(_reap.os, 'unlink', unlink)

    _reap.unlink_uds_bind_addrs(
        proc=proc,
        subactor=subactor,
    )

    unlink.assert_not_called()


def test_uds_bindspace_created_implicitly(
    debug_mode: bool,
    bindspace_dir_str: str,
):
    registry_addr: tuple = (
        f'{bindspace_dir_str}',
        'registry@doggy.sock',
    )
    bs_dir_str: str = registry_addr[0]

    # XXX, ensure bindspace-dir DNE beforehand!
    assert not Path(bs_dir_str).is_dir()

    async def main():
        async with tractor.open_nursery(
            enable_transports=['uds'],
            registry_addrs=[registry_addr],
            debug_mode=debug_mode,
        ) as _an:

            # XXX MUST be created implicitly by
            # `.ipc._uds.start_listener()`!
            assert Path(bs_dir_str).is_dir()

            root: Actor = tractor.current_actor()
            assert root.is_registrar

            assert registry_addr in root.reg_addrs
            assert (
                registry_addr
                in
                _state._runtime_vars['_registry_addrs']
            )
            assert (
                _addr.wrap_address(registry_addr)
                in
                root.registry_addrs
            )

    trio.run(main)


def test_uds_double_listen_raises_connerr(
    debug_mode: bool,
    bindspace_dir_str: str,
):
    registry_addr: tuple = (
        f'{bindspace_dir_str}',
        'registry@doggy.sock',
    )

    async def main():
        async with tractor.open_nursery(
            enable_transports=['uds'],
            registry_addrs=[registry_addr],
            debug_mode=debug_mode,
        ) as _an:

            # runtime up
            root: Actor = tractor.current_actor()

            from tractor.ipc._uds import (
                start_listener,
                UDSAddress,
            )
            ya_bound_addr: UDSAddress = root.registry_addrs[0]
            try:
                await start_listener(
                    addr=ya_bound_addr,
                )
            except ConnectionError as connerr:
                assert type(src_exc := connerr.__context__) is OSError
                assert 'Address already in use' in src_exc.args
                # complete, exit test.

            else:
                pytest.fail('It dint raise a connerr !?')


    trio.run(main)
