"""Tests for Stream and Event semantics.

Test coverage:
- Event lifecycle (pending -> recorded -> completed)
- Event wait() with timeout
- Event is_complete()
- Stream record_event()
- Stream wait_event() with dependencies
- Stream synchronize()
- Error cases: repeat wait, invalid event types, timeout
"""

from __future__ import annotations

import time

import pytest

from sagellm_comm.overlap import CPUEventHandle, EventState, create_cpu_stream


class TestCPUEventHandle:
    """Tests for CPUEventHandle semantics."""

    def test_event_initial_state(self) -> None:
        """Event should start in PENDING state."""
        event = CPUEventHandle()
        assert event.state == EventState.PENDING
        assert event.recorded_at is None
        assert event.completed_at is None
        assert not event.is_complete()

    def test_event_record_and_complete(self) -> None:
        """Event should transition PENDING -> RECORDED -> COMPLETED."""
        event = CPUEventHandle()

        # Mark recorded
        event._mark_recorded()
        assert event.state == EventState.RECORDED
        assert event.recorded_at is not None
        assert event.completed_at is None
        assert not event.is_complete()

        # Mark completed
        event._mark_completed()
        assert event.state == EventState.COMPLETED
        assert event.completed_at is not None
        assert event.is_complete()

    def test_event_wait_already_complete(self) -> None:
        """wait() on completed event should return immediately."""
        event = CPUEventHandle()
        event._mark_recorded()
        event._mark_completed()

        start = time.time()
        event.wait()
        elapsed = time.time() - start

        assert elapsed < 0.01  # Should be nearly instant
        assert event.is_complete()

    def test_event_wait_with_completion(self) -> None:
        """wait() should block until event is completed."""
        event = CPUEventHandle()
        event._mark_recorded()

        # Complete event in background after delay
        import threading

        def _complete_later() -> None:
            time.sleep(0.05)  # 50ms delay
            event._mark_completed()

        threading.Thread(target=_complete_later, daemon=True).start()

        start = time.time()
        event.wait()
        elapsed = time.time() - start

        assert elapsed >= 0.05  # Should have waited
        assert event.is_complete()

    def test_event_wait_timeout(self) -> None:
        """wait() should raise TimeoutError if timeout expires."""
        event = CPUEventHandle()
        event._mark_recorded()

        with pytest.raises(TimeoutError, match="did not complete within"):
            event.wait(timeout=0.05)  # 50ms timeout

        assert not event.is_complete()

    def test_event_wait_error_state(self) -> None:
        """wait() should raise RuntimeError if event is in ERROR state."""
        event = CPUEventHandle()
        event._mark_recorded()
        event._mark_error()

        with pytest.raises(RuntimeError, match="is in ERROR state"):
            event.wait()

    def test_event_repeat_wait_idempotent(self) -> None:
        """Multiple wait() calls on same event should be idempotent."""
        event = CPUEventHandle()
        event._mark_recorded()
        event._mark_completed()

        # First wait
        event.wait()
        assert event.is_complete()

        # Second wait (should not raise)
        event.wait()
        assert event.is_complete()

    def test_event_elapsed_ms(self) -> None:
        """elapsed_ms() should return correct elapsed time."""
        event = CPUEventHandle()
        event._mark_recorded()

        # Sleep a bit
        time.sleep(0.05)

        event._mark_completed()

        elapsed = event.elapsed_ms()
        assert elapsed is not None
        assert elapsed >= 50.0  # At least 50ms
        assert elapsed < 100.0  # But not too much more


class TestCPUStreamHandle:
    """Tests for CPUStreamHandle semantics."""

    def test_stream_initial_state(self) -> None:
        """Stream should be idle initially."""
        stream = create_cpu_stream()
        assert stream.is_idle()
        assert stream.stream_id is not None
        assert stream.trace_id is not None

    def test_stream_record_event_immediate_completion(self) -> None:
        """record_event() on idle stream should complete immediately."""
        stream = create_cpu_stream()

        event = stream.record_event()
        assert isinstance(event, CPUEventHandle)
        assert event.state == EventState.COMPLETED  # Idle stream -> immediate completion
        assert event.is_complete()

    def test_stream_record_event_with_pending_ops(self) -> None:
        """record_event() with pending ops should not complete immediately."""
        stream = create_cpu_stream()

        # Submit operation
        stream._submit_operation()
        assert not stream.is_idle()

        # Record event
        event = stream.record_event()
        assert event.state == EventState.RECORDED
        assert not event.is_complete()

        # Complete operation
        stream._complete_operation()
        assert stream.is_idle()
        assert event.is_complete()

    def test_stream_wait_event_with_dependency(self) -> None:
        """wait_event() should block until dependency is satisfied."""
        stream1 = create_cpu_stream()
        stream2 = create_cpu_stream()

        # Submit operation on stream1
        stream1._submit_operation()
        event1 = stream1.record_event()
        assert not event1.is_complete()

        # stream2 waits for event1 in background
        import threading

        wait_started = threading.Event()
        wait_completed = threading.Event()

        def _wait_for_event() -> None:
            wait_started.set()
            stream2.wait_event(event1)
            wait_completed.set()

        wait_thread = threading.Thread(target=_wait_for_event, daemon=True)
        wait_thread.start()

        # Wait for thread to start waiting
        wait_started.wait(timeout=1.0)
        time.sleep(0.01)  # Give it time to block

        # Should not have completed yet
        assert not wait_completed.is_set()

        # Complete stream1 operation
        stream1._complete_operation()

        # Now wait should complete
        wait_completed.wait(timeout=1.0)
        assert wait_completed.is_set()
        assert event1.is_complete()

    def test_stream_wait_event_invalid_type(self) -> None:
        """wait_event() should reject non-CPUEventHandle."""
        stream = create_cpu_stream()

        # Create a fake event from different implementation
        class FakeEvent:
            pass

        with pytest.raises(ValueError, match="requires CPUEventHandle"):
            stream.wait_event(FakeEvent())  # type: ignore

    def test_stream_synchronize(self) -> None:
        """synchronize() should wait for all pending operations."""
        stream = create_cpu_stream()

        # Submit operation
        stream._submit_operation()
        assert not stream.is_idle()

        # Complete operation in background
        import threading

        def _complete_later() -> None:
            time.sleep(0.05)
            stream._complete_operation()

        threading.Thread(target=_complete_later, daemon=True).start()

        # Synchronize (should block)
        start = time.time()
        stream.synchronize()
        elapsed = time.time() - start

        assert elapsed >= 0.05
        assert stream.is_idle()

    def test_stream_synchronize_completes_events(self) -> None:
        """synchronize() should complete all recorded events."""
        stream = create_cpu_stream()

        # Submit multiple operations
        stream._submit_operation()
        event1 = stream.record_event()
        stream._submit_operation()
        event2 = stream.record_event()

        assert not event1.is_complete()
        assert not event2.is_complete()

        # Complete operations
        stream._complete_operation()
        stream._complete_operation()

        # Synchronize
        stream.synchronize()

        assert event1.is_complete()
        assert event2.is_complete()

    def test_stream_operations_tracking(self) -> None:
        """_submit_operation and _complete_operation should track pending ops."""
        stream = create_cpu_stream()

        assert stream.is_idle()

        stream._submit_operation()
        assert not stream.is_idle()

        stream._submit_operation()
        assert not stream.is_idle()

        stream._complete_operation()
        assert not stream.is_idle()

        stream._complete_operation()
        assert stream.is_idle()


class TestStreamEventIntegration:
    """Integration tests for stream-event interaction."""

    def test_two_stream_overlap(self) -> None:
        """Test basic 2-stream overlap scenario."""
        comm_stream = create_cpu_stream(trace_id="test-trace")
        compute_stream = create_cpu_stream(trace_id="test-trace")

        # Simulate communication
        comm_stream._submit_operation()

        # Record event
        comm_event = comm_stream.record_event()
        assert not comm_event.is_complete()

        # Compute stream waits for comm event in background
        import threading

        wait_completed = threading.Event()

        def _compute_after_comm() -> None:
            compute_stream.wait_event(comm_event)
            compute_stream._submit_operation()
            wait_completed.set()

        compute_thread = threading.Thread(target=_compute_after_comm, daemon=True)
        compute_thread.start()

        # Complete communication
        time.sleep(0.02)
        comm_stream._complete_operation()

        # Wait should complete
        wait_completed.wait(timeout=1.0)
        assert wait_completed.is_set()
        assert comm_event.is_complete()

        # Cleanup
        compute_stream._complete_operation()
        compute_stream.synchronize()

    def test_trace_id_propagation(self) -> None:
        """Events should inherit trace_id from stream."""
        trace_id = "test-trace-12345"
        stream = create_cpu_stream(trace_id=trace_id)

        event = stream.record_event()

        assert stream.trace_id == trace_id
        assert event.trace_id == trace_id
