# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import multiprocessing as mp

import pytest
from helpers.buffers import PatternGen
from helpers.child_processes import child_timeout_sec, kill_subprocesses
from helpers.constants import POOL_SIZE

from cuda.core import Device, DeviceMemoryResource, DeviceMemoryResourceOptions
from cuda.core._utils.cuda_utils import CUDAError

CHILD_TIMEOUT_SEC = child_timeout_sec()
NBYTES = 64

# these tests spawn new processes and files which fails for very many threads
pytestmark = pytest.mark.parallel_threads_limit(4)


class TestPeerAccessNotPreservedOnImport:
    """
    Verify that peer access settings are not preserved when a memory resource
    is sent to another process via IPC, and that peer access can be set after import.
    """

    @pytest.mark.flaky(reruns=2)
    def test_main(self, ipc_mempool_device_x2):
        dev0, dev1 = ipc_mempool_device_x2

        # Parent Process - Create and Configure MR
        dev1.set_current()
        options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
        mr = DeviceMemoryResource(dev1, options=options)
        mr.peer_accessible_by = [dev0]
        assert mr.peer_accessible_by == {dev0}

        # Spawn child process
        process = mp.Process(target=self.child_main, args=(mr,))
        process.start()
        process.join(timeout=CHILD_TIMEOUT_SEC)
        survivors = kill_subprocesses(process)
        assert not survivors, "child did not exit within timeout"
        assert process.exitcode == 0

        # Verify parent's MR still has peer access set (independent state)
        assert mr.peer_accessible_by == {dev0}
        mr.close()

    def child_main(self, mr):
        Device(1).set_current()
        assert mr.is_mapped is True
        assert mr.device_id == 1
        assert mr.peer_accessible_by == set()
        mr.peer_accessible_by = [0]
        assert mr.peer_accessible_by == {Device(0)}
        mr.peer_accessible_by = []
        assert mr.peer_accessible_by == set()
        mr.close()


class TestBufferPeerAccessAfterImport:
    """
    Verify that buffers imported via IPC can be accessed from peer devices after
    setting peer access on the imported memory resource, and that access can be revoked.
    """

    @pytest.mark.flaky(reruns=2)
    @pytest.mark.parametrize("grant_access_in_parent", [True, False])
    def test_main(self, ipc_mempool_device_x2, grant_access_in_parent):
        dev0, dev1 = ipc_mempool_device_x2

        # Parent Process - Create MR and Buffer
        dev1.set_current()
        options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
        mr = DeviceMemoryResource(dev1, options=options)
        if grant_access_in_parent:
            mr.peer_accessible_by = [dev0]
            assert mr.peer_accessible_by == {dev0}
        else:
            assert mr.peer_accessible_by == set()
        buffer = mr.allocate(NBYTES, stream=dev1.default_stream)
        pgen = PatternGen(dev1, NBYTES)
        pgen.fill_buffer(buffer, seed=False)

        # Spawn child process
        process = mp.Process(target=self.child_main, args=(mr, buffer))
        process.start()
        process.join(timeout=CHILD_TIMEOUT_SEC)
        survivors = kill_subprocesses(process)
        assert not survivors, "child did not exit within timeout"
        assert process.exitcode == 0

        buffer.close()
        # TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
        dev1.sync()
        mr.close()

    def child_main(self, mr, buffer):
        # Verify MR and buffer are mapped
        Device(1).set_current()
        assert mr.is_mapped is True
        assert buffer.is_mapped is True
        assert mr.device_id == 1
        assert buffer.device_id == 1

        # Test 1: Buffer accessible from resident device (dev1) - should always work
        dev1 = Device(1)
        dev1.set_current()
        PatternGen(dev1, NBYTES).verify_buffer(buffer, seed=False)

        # Test 2: Buffer NOT accessible from dev0 initially (peer access not preserved)
        dev0 = Device(0)
        dev0.set_current()
        with pytest.raises(CUDAError, match="CUDA_ERROR_INVALID_VALUE"):
            PatternGen(dev0, NBYTES).verify_buffer(buffer, seed=False)

        # Test 3: Set peer access and verify buffer becomes accessible
        dev1.set_current()
        mr.peer_accessible_by = [0]
        assert mr.peer_accessible_by == {dev0}
        dev0.set_current()
        PatternGen(dev0, NBYTES).verify_buffer(buffer, seed=False)

        # Test 4: Revoke peer access and verify buffer becomes inaccessible
        dev1.set_current()
        mr.peer_accessible_by = []
        assert mr.peer_accessible_by == set()
        dev0.set_current()
        with pytest.raises(CUDAError, match="CUDA_ERROR_INVALID_VALUE"):
            PatternGen(dev0, NBYTES).verify_buffer(buffer, seed=False)

        buffer.close()
        # TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
        dev1.sync()
        mr.close()
