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

import pytest
from helpers.misc import StreamWrapper

from cuda.core import Device, Stream, StreamOptions
from cuda.core._event import Event
from cuda.core._stream import LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM, Stream_accept
from cuda.core._utils.cuda_utils import driver


def test_stream_init_disabled():
    with pytest.raises(RuntimeError, match=r"^Stream objects cannot be instantiated directly\."):
        Stream()  # Reject at front door.


def test_stream_accept_rejects_none():
    """Stream_accept(None) raises TypeError so APIs cannot silently fall back
    to the default stream (issue #2001)."""
    with pytest.raises(TypeError, match=r"stream is required and must not be None"):
        Stream_accept(None)


def test_stream_init_with_options(init_cuda):
    stream = Device().create_stream(options=StreamOptions(nonblocking=True, priority=0))
    assert stream.is_nonblocking is True
    assert stream.priority == 0


def test_stream_handle(init_cuda):
    stream = Device().create_stream(options=StreamOptions())
    assert isinstance(stream.handle, driver.CUstream)


def test_stream_is_nonblocking(init_cuda):
    stream = Device().create_stream(options=StreamOptions(nonblocking=True))
    assert stream.is_nonblocking is True


def test_stream_priority(init_cuda):
    stream = Device().create_stream(options=StreamOptions(priority=0))
    assert stream.priority == 0
    stream = Device().create_stream(options=StreamOptions(priority=-1))
    assert stream.priority == -1
    with pytest.raises(ValueError):
        stream = Device().create_stream(options=StreamOptions(priority=1))


def test_stream_sync(init_cuda):
    stream = Device().create_stream(options=StreamOptions())
    stream.sync()  # Should not raise any exceptions


def test_stream_record(init_cuda):
    stream = Device().create_stream(options=StreamOptions())
    event = stream.record()
    assert isinstance(event, Event)


def test_stream_wait_event(init_cuda):
    s1 = Device().create_stream()
    s2 = Device().create_stream()
    e1 = s1.record()
    s2.wait(e1)  # Should not raise any exceptions
    s2.sync()


def test_stream_wait_invalid_event(init_cuda):
    stream = Device().create_stream(options=StreamOptions())
    with pytest.raises(ValueError):
        stream.wait(event_or_stream="invalid_event")


def test_stream_device(init_cuda):
    stream = Device().create_stream(options=StreamOptions())
    device = stream.device
    assert isinstance(device, Device)


def test_stream_context(init_cuda):
    stream = Device().create_stream(options=StreamOptions())
    context = stream.context
    assert context is not None
    assert context.handle is not None


def test_stream_from_foreign_stream(init_cuda):
    device = Device()
    other_stream = device.create_stream(options=StreamOptions())
    stream = device.create_stream(obj=StreamWrapper(other_stream))
    # Now that __eq__ is implemented (issue #664), we can compare directly
    assert other_stream == stream
    assert hash(other_stream) == hash(stream)
    device = stream.device
    assert isinstance(device, Device)
    context = stream.context
    assert context is not None


def test_stream_from_handle():
    stream = Stream.from_handle(0)
    assert isinstance(stream, Stream)


def test_legacy_default_stream():
    assert isinstance(LEGACY_DEFAULT_STREAM, Stream)


def test_per_thread_default_stream():
    assert isinstance(PER_THREAD_DEFAULT_STREAM, Stream)


def test_stream_subclassing(init_cuda):
    class MyStream(Stream):
        pass

    dev = Device()
    dev.set_current()
    stream = MyStream._init(options=StreamOptions(), device_id=dev.device_id)
    assert isinstance(stream, MyStream)


def test_stream_legacy_default_subclassing():
    class MyStream(Stream):
        pass

    stream = MyStream._legacy_default()
    assert isinstance(stream, MyStream)


def test_stream_per_thread_default_subclassing():
    class MyStream(Stream):
        pass

    stream = MyStream._per_thread_default()
    assert isinstance(stream, MyStream)


# ============================================================================
# Stream Equality Tests
# ============================================================================


def test_stream_equality_same_handle(init_cuda):
    """Two Stream objects wrapping same handle should be equal."""
    device = Device()
    s1 = device.create_stream()
    s2 = Stream.from_handle(int(s1.handle))

    assert s1 == s2, "Equal streams should be equal"


def test_stream_inequality_different_handles(init_cuda):
    """Different streams should not be equal."""
    device = Device()
    s1 = device.create_stream()
    s2 = device.create_stream()

    assert s1 != s2, "Different streams should not be equal"


def test_stream_equality_reflexive(init_cuda):
    """Stream should equal itself (reflexive property)."""
    stream = Device().create_stream()
    assert stream == stream, "Stream should equal itself"


def test_stream_equality_symmetric(init_cuda):
    """If s1 == s2, then s2 == s1 (symmetric property)."""
    device = Device()
    s1 = device.create_stream()
    s2 = Stream.from_handle(int(s1.handle))

    assert s1 == s2, "Equality should be symmetric"
    assert s2 == s1, "Equality should be symmetric"


def test_stream_type_safety(init_cuda):
    """Comparing Stream with wrong type should return False."""
    stream = Device().create_stream()

    # These should not raise exceptions
    assert (stream == "not a stream") is False
    assert (stream == 123) is False
    assert (stream is None) is False
    assert (stream == Device()) is False


def test_stream_not_equal_operator(init_cuda):
    """Test != operator works correctly for streams."""
    device = Device()
    s1 = device.create_stream()
    s2 = device.create_stream()
    s3 = Stream.from_handle(int(s1.handle))

    assert s1 != s2, "Different streams should be not-equal"
    assert s1 == s3, "Same handle streams should be equal"


# ============================================================================
# Stream Hash Tests
# ============================================================================


def test_stream_hash_consistency(init_cuda):
    """Hash of same Stream object should be consistent."""
    stream = Device().create_stream()
    hash1 = hash(stream)
    hash2 = hash(stream)
    assert hash1 == hash2, "Hash should be consistent for same object"


def test_stream_hash_equality_same_handle(init_cuda):
    """Two Stream objects wrapping same handle should hash equal."""
    device = Device()
    s1 = device.create_stream()
    s2 = Stream.from_handle(int(s1.handle))

    assert hash(s1) == hash(s2), "Same handle should produce same hash"
    assert s1 == s2, "Streams with same handle should be equal"


def test_stream_hash_inequality_different_handles(init_cuda):
    """Different streams should have different hashes."""
    device = Device()
    s1 = device.create_stream()
    s2 = device.create_stream()

    assert s1 != s2, "Different streams should not be equal"
    # Hash inequality not guaranteed by contract, but very likely
    assert hash(s1) != hash(s2), "Different streams likely have different hashes"


def test_stream_dict_key(init_cuda):
    """Streams should be usable as dictionary keys."""
    device = Device()
    s1 = device.create_stream()
    s2 = Stream.from_handle(int(s1.handle))

    # Use s1 as key
    cache = {s1: "value1"}

    # s2 should find s1's entry (same handle)
    assert s2 in cache, "Stream with same handle should be in dict"
    assert cache[s2] == "value1", "Should retrieve value via equivalent stream"

    # Updating via s2 should update s1's entry
    cache[s2] = "value2"
    assert len(cache) == 1, "Should update existing entry, not add new one"
    assert cache[s1] == "value2", "Original key should see updated value"


def test_stream_set_membership(init_cuda):
    """Streams should work correctly in sets."""
    device = Device()
    s1 = device.create_stream()
    s2 = Stream.from_handle(int(s1.handle))
    s3 = device.create_stream()

    stream_set = {s1}

    # s2 wraps same handle as s1, should be recognized as duplicate
    assert s2 in stream_set, "Equivalent stream should be in set"
    stream_set.add(s2)
    assert len(stream_set) == 1, "Should not add duplicate"

    # s3 is different stream
    stream_set.add(s3)
    assert len(stream_set) == 2, "Should add different stream"


def test_builtin_stream_hash(init_cuda):
    """Builtin streams should be hashable and usable as dict keys."""
    Device().set_current()

    legacy = LEGACY_DEFAULT_STREAM
    per_thread = PER_THREAD_DEFAULT_STREAM

    # Should be hashable
    hash_legacy = hash(legacy)
    hash_per_thread = hash(per_thread)
    assert isinstance(hash_legacy, int)
    assert isinstance(hash_per_thread, int)

    # Should have different hashes (different handles)
    assert hash_legacy != hash_per_thread

    # Should be usable in dicts
    cache = {legacy: "legacy_data", per_thread: "per_thread_data"}
    assert len(cache) == 2
    assert cache[legacy] == "legacy_data"
    assert cache[per_thread] == "per_thread_data"


def test_default_stream_consistency(init_cuda):
    """Device.default_stream should be consistent and hashable."""
    device = Device()

    default1 = device.default_stream
    default2 = device.default_stream

    # Should be same object (or at least equal)
    assert default1 == default2
    assert hash(default1) == hash(default2)


class _BadStreamProtocol:
    """Object whose __cuda_stream__ (a method) returns a malformed value."""

    def __init__(self, value):
        self._value = value

    def __cuda_stream__(self):
        return self._value


class _AttrStreamProtocol:
    """Object implementing __cuda_stream__ as an attribute (deprecated form)
    rather than a method; the tuple length is wrong so resolution stops before
    any GPU work."""

    __cuda_stream__ = (0, 1, 2)


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_stream_init_rejects_obj_and_options():
    """Stream._init rejects supplying both a foreign object and options."""
    from cuda.core._stream import Stream

    with pytest.raises(ValueError, match="obj and options cannot be both specified"):
        Stream._init(obj=_BadStreamProtocol((0, 0)), options=StreamOptions())


@pytest.mark.agent_authored(model="claude-opus-4.8")
@pytest.mark.parametrize(
    "value,match",
    [
        ((0, 1, 2), "must return a sequence with 2 elements"),  # wrong length
        (5, "must return a sequence with 2 elements"),  # not a sequence
        ((1, 123), r"first element of the sequence.*must be 0"),  # bad version
    ],
)
def test_stream_init_rejects_bad_cuda_stream_protocol(value, match):
    """A foreign object whose __cuda_stream__ returns a malformed value is
    rejected before any handle is created."""
    from cuda.core._stream import Stream

    with pytest.raises(RuntimeError, match=match):
        Stream._init(obj=_BadStreamProtocol(value))


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_stream_init_warns_on_attribute_cuda_stream_protocol():
    """Implementing __cuda_stream__ as an attribute (not a method) is deprecated:
    resolution emits a DeprecationWarning and then still rejects the malformed
    (wrong-length) value with a RuntimeError."""
    from cuda.core._stream import Stream

    with (
        pytest.warns(DeprecationWarning, match="must be implemented as a method"),
        pytest.raises(RuntimeError, match="must return a sequence with 2 elements"),
    ):
        Stream._init(obj=_AttrStreamProtocol())


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_stream_init_from_existing_stream_object(init_cuda):
    """Passing an existing Stream as the foreign object yields a borrowed stream over the same handle."""
    from cuda.core._stream import Stream

    src = Device().create_stream(options=StreamOptions())
    borrowed = Stream._init(obj=src)
    assert int(borrowed.handle) == int(src.handle)
    borrowed.close()
    src.close()


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_stream_from_handle_lazy_flag_and_priority_queries(init_cuda):
    """A from_handle stream reports is_nonblocking and priority read back from the
    driver (matching the source stream), not the constructor defaults."""
    from cuda.core._stream import Stream

    # priority=-1 (not the default 0) so the value proves the driver was actually queried.
    real = Device().create_stream(options=StreamOptions(nonblocking=True, priority=-1))
    wrapped = Stream.from_handle(int(real.handle))
    assert wrapped.is_nonblocking is True
    assert wrapped.priority == -1
    wrapped.close()
    real.close()


@pytest.mark.agent_authored(model="claude-opus-4.8")
@pytest.mark.thread_unsafe(
    reason="mutates the process-global CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM env var that default_stream() reads live"
)
def test_default_stream_per_thread_when_env_set(monkeypatch):
    """default_stream() returns the per-thread default stream when
    CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM is set to a nonzero value, and the
    legacy default stream otherwise."""
    from cuda.core._stream import default_stream

    monkeypatch.setenv("CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM", "1")
    assert default_stream() is PER_THREAD_DEFAULT_STREAM
    monkeypatch.delenv("CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM", raising=False)
    assert default_stream() is LEGACY_DEFAULT_STREAM


def _skip_unless_multi_gpu():
    from cuda.core import system

    if system.get_num_devices() < 2:
        pytest.skip("requires 2+ GPUs")


@pytest.mark.agent_authored(model="claude-opus-5")
@pytest.mark.parametrize("stream", [LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM])
def test_default_stream_follows_current_context(stream):
    """A default-stream token denotes whatever context is current, so its
    queries follow a context switch instead of reporting the first context
    they ever saw (issue #2485)."""
    _skip_unless_multi_gpu()

    for device_id in (0, 1, 0):
        dev = Device(device_id)
        dev.set_current()
        assert stream.device.device_id == device_id
        assert stream.context == dev.context
        assert stream.record().context == dev.context
        # Exercise Stream.resources and check it tracks the same context
        # resolution as .context (issue #2485).
        assert stream.resources.sm.sm_count == stream.context.resources.sm.sm_count
        assert stream.resources.sm.sm_count == dev.resources.sm.sm_count


@pytest.mark.agent_authored(model="claude-opus-5")
def test_default_stream_first_touch_does_not_pin_context():
    """Any query resolves the context, including repr(), so logging a default
    stream must not bind the singleton to whichever context happened to be
    current at the time (issue #2485)."""
    _skip_unless_multi_gpu()

    Device(1).set_current()
    repr(LEGACY_DEFAULT_STREAM)

    dev0 = Device(0)
    dev0.set_current()
    assert LEGACY_DEFAULT_STREAM.device.device_id == 0
    assert LEGACY_DEFAULT_STREAM.context == dev0.context


@pytest.mark.agent_authored(model="claude-opus-5")
def test_created_stream_keeps_its_own_context():
    """A created stream has a context fixed at creation and keeps reporting it
    across a switch; only default-stream tokens follow the current context
    (issue #2485)."""
    _skip_unless_multi_gpu()

    dev0 = Device(0)
    dev0.set_current()
    stream = dev0.create_stream()
    assert stream.context == dev0.context

    try:
        Device(1).set_current()
        assert stream.device.device_id == 0
        assert stream.context == dev0.context
    finally:
        dev0.set_current()
        stream.close()
