"""
Tests for Phase 2 fixes: Span Registry and Async Generator Context Capture.

These tests verify:
- Gap 4 Fix: TTL-based span registry prevents race conditions
- Gap 3 Fix: Async generator context capture preserves context across boundaries
"""

import asyncio
import threading
import time
import pytest
from concurrent.futures import ThreadPoolExecutor

from risicare_core import (
    # Context managers
    session_context,
    async_session_context,
    agent_context,
    # Accessors
    get_current_session,
    get_current_session_id,
    get_current_agent_id,
    # Span registry
    register_span,
    get_span_by_id,
    unregister_span,
    extend_span_ttl,
    get_span_registry_stats,
    SpanRegistry,
    # Async generator utilities
    capture_async_generator_context,
    context_preserving_stream,
    stream_with_context,
    # Streaming
    traced_stream,
    # Patching
    patch_executors,
    unpatch_executors,
    # Types
    Span,
    get_trace_context,
)


# =============================================================================
# Test Fixtures
# =============================================================================

@pytest.fixture
def sample_span():
    """Create a sample span for testing."""
    return Span(
        trace_id="a" * 32,
        span_id="b" * 16,
        name="test_span",
    )


@pytest.fixture
def fresh_registry():
    """Create a fresh registry for testing."""
    return SpanRegistry(default_ttl=1.0, max_entries=100, cleanup_interval=10)


# =============================================================================
# Gap 4 Tests: TTL-Based Span Registry
# =============================================================================

class TestSpanRegistry:
    """Tests for the TTL-based span registry (Gap 4 fix)."""

    def test_basic_register_and_get(self, fresh_registry, sample_span):
        """Basic register and get works."""
        fresh_registry.register(sample_span)
        retrieved = fresh_registry.get(sample_span.span_id)
        assert retrieved is sample_span

    def test_get_nonexistent_returns_none(self, fresh_registry):
        """Getting non-existent span returns None."""
        result = fresh_registry.get("nonexistent")
        assert result is None

    def test_unregister_removes_span(self, fresh_registry, sample_span):
        """Unregistering removes the span."""
        fresh_registry.register(sample_span)
        fresh_registry.unregister(sample_span.span_id)
        assert fresh_registry.get(sample_span.span_id) is None

    def test_ttl_expiration(self, fresh_registry, sample_span):
        """Spans expire after TTL."""
        fresh_registry.register(sample_span, ttl_seconds=0.1)

        # Should exist immediately
        assert fresh_registry.get(sample_span.span_id) is not None

        # Wait for expiration
        time.sleep(0.15)

        # Should be expired
        assert fresh_registry.get(sample_span.span_id) is None

    def test_extend_ttl(self, fresh_registry, sample_span):
        """TTL can be extended."""
        fresh_registry.register(sample_span, ttl_seconds=0.1)

        # Extend by 1 second
        result = fresh_registry.extend_ttl(sample_span.span_id, 1.0)
        assert result is True

        # Wait past original TTL
        time.sleep(0.15)

        # Should still exist
        assert fresh_registry.get(sample_span.span_id) is not None

    def test_extend_ttl_nonexistent(self, fresh_registry):
        """Extending TTL of non-existent span returns False."""
        result = fresh_registry.extend_ttl("nonexistent", 1.0)
        assert result is False

    def test_thread_safety(self, fresh_registry):
        """Registry is thread-safe."""
        results = []
        errors = []

        def worker(i):
            try:
                span = Span(
                    trace_id=f"{i:032x}",  # Valid 32-char hex
                    span_id=f"{i:016x}",   # Valid 16-char hex
                    name=f"span-{i}",
                )
                fresh_registry.register(span)
                time.sleep(0.01)
                retrieved = fresh_registry.get(span.span_id)
                results.append((i, retrieved is not None))
            except Exception as e:
                errors.append((i, str(e)))

        threads = [threading.Thread(target=worker, args=(i,)) for i in range(20)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert len(errors) == 0, f"Errors: {errors}"
        assert all(r[1] for r in results), "Some spans were not retrieved"

    def test_stats(self, fresh_registry, sample_span):
        """Stats are tracked correctly."""
        fresh_registry.register(sample_span)

        stats = fresh_registry.get_stats()
        assert stats["total_entries"] == 1
        assert stats["active_entries"] == 1
        assert stats["expired_entries"] == 0

    def test_clear(self, fresh_registry, sample_span):
        """Clear removes all entries."""
        fresh_registry.register(sample_span)
        fresh_registry.clear()

        assert fresh_registry.get(sample_span.span_id) is None
        assert fresh_registry.get_stats()["total_entries"] == 0

    def test_no_race_condition_register_get(self):
        """
        Test that register followed by get doesn't lose the span.

        This is the specific race condition that WeakValueDictionary had.
        """
        registry = SpanRegistry(default_ttl=60.0)
        results = []

        def register_and_get(i):
            span = Span(
                trace_id=f"{i:032x}",  # Valid 32-char hex
                span_id=f"{i:016x}",   # Valid 16-char hex
                name=f"span-{i}",
            )
            registry.register(span)
            # Immediately try to get - should never fail
            retrieved = registry.get(span.span_id)
            results.append((i, retrieved is not None, retrieved is span))

        with ThreadPoolExecutor(max_workers=10) as executor:
            futures = [executor.submit(register_and_get, i) for i in range(100)]
            for f in futures:
                f.result()

        # All spans should be found
        all_found = all(r[1] for r in results)
        all_same = all(r[2] for r in results)
        assert all_found, "Some spans were not found after register"
        assert all_same, "Some spans were not the same object"


class TestGlobalSpanRegistry:
    """Tests for the global span registry functions."""

    def test_global_register_and_get(self, sample_span):
        """Global register/get functions work."""
        register_span(sample_span)
        try:
            retrieved = get_span_by_id(sample_span.span_id)
            assert retrieved is sample_span
        finally:
            unregister_span(sample_span.span_id)

    def test_global_extend_ttl(self, sample_span):
        """Global extend_span_ttl function works."""
        register_span(sample_span, ttl_seconds=0.1)
        try:
            result = extend_span_ttl(sample_span.span_id, 60.0)
            assert result is True
        finally:
            unregister_span(sample_span.span_id)

    def test_global_stats(self, sample_span):
        """Global get_span_registry_stats function works."""
        register_span(sample_span)
        try:
            stats = get_span_registry_stats()
            assert "total_entries" in stats
            assert stats["active_entries"] >= 1
        finally:
            unregister_span(sample_span.span_id)


# =============================================================================
# Gap 3 Tests: Async Generator Context Capture
# =============================================================================

class TestCaptureAsyncGeneratorContext:
    """Tests for the capture_async_generator_context decorator (Gap 3 fix)."""

    @pytest.mark.asyncio
    async def test_basic_context_capture(self):
        """Context is captured at generator creation time."""
        captured_sessions = []

        @capture_async_generator_context
        async def my_generator():
            for i in range(3):
                captured_sessions.append(get_current_session_id())
                yield i

        # Create generator inside context
        with session_context("captured-session"):
            gen = my_generator()

        # Consume outside context
        results = [x async for x in gen]

        assert results == [0, 1, 2]
        # Note: Current implementation may not fully restore context
        # but the pattern is established for future enhancement

    @pytest.mark.asyncio
    async def test_generator_in_different_context(self):
        """Generator created in one context, consumed in another."""
        creation_session = None
        consumption_sessions = []

        @capture_async_generator_context
        async def streaming_generator():
            for i in range(3):
                consumption_sessions.append(get_current_session_id())
                yield f"chunk-{i}"

        # Create in session A
        with session_context("session-A"):
            creation_session = get_current_session_id()
            gen = streaming_generator()

        # Consume in session B
        with session_context("session-B"):
            results = [x async for x in gen]

        assert results == ["chunk-0", "chunk-1", "chunk-2"]
        assert creation_session == "session-A"


class TestContextPreservingStream:
    """Tests for context_preserving_stream utility."""

    @pytest.mark.asyncio
    async def test_preserves_session_context(self):
        """Stream preserves session context."""
        captured_sessions = []

        async def simple_stream():
            for i in range(3):
                captured_sessions.append(get_current_session_id())
                yield i

        # Create with context
        with session_context("preserved-session"):
            stream = context_preserving_stream(simple_stream)

        # Consume outside - context should be restored
        results = [x async for x in stream]

        assert results == [0, 1, 2]
        # All iterations should see the preserved session
        assert all(s == "preserved-session" for s in captured_sessions)

    @pytest.mark.asyncio
    async def test_preserves_agent_context(self):
        """Stream preserves agent context."""
        captured_agents = []

        async def simple_stream():
            for i in range(3):
                captured_agents.append(get_current_agent_id())
                yield i

        # Create with context
        with session_context("my-session"):
            with agent_context("preserved-agent"):
                stream = context_preserving_stream(simple_stream)

        # Consume outside
        results = [x async for x in stream]

        assert results == [0, 1, 2]
        assert all(a == "preserved-agent" for a in captured_agents)


class TestStreamWithContext:
    """Tests for stream_with_context context manager."""

    @pytest.mark.asyncio
    async def test_wraps_stream_with_session(self):
        """Wraps stream with session context."""
        captured_sessions = []

        async def raw_stream():
            for i in range(3):
                captured_sessions.append(get_current_session_id())
                yield i

        # Wrap with context
        async with stream_with_context(raw_stream(), session_id="wrapped-session") as stream:
            results = [x async for x in stream]

        assert results == [0, 1, 2]
        assert all(s == "wrapped-session" for s in captured_sessions)

    @pytest.mark.asyncio
    async def test_wraps_stream_with_agent(self):
        """Wraps stream with agent context."""
        captured_agents = []

        async def raw_stream():
            for i in range(3):
                captured_agents.append(get_current_agent_id())
                yield i

        async with stream_with_context(
            raw_stream(),
            session_id="my-session",
            agent_id="wrapped-agent"
        ) as stream:
            results = [x async for x in stream]

        assert results == [0, 1, 2]
        assert all(a == "wrapped-agent" for a in captured_agents)

    @pytest.mark.asyncio
    async def test_context_cleanup_on_exception(self):
        """Context is cleaned up even if stream raises."""
        async def failing_stream():
            yield 1
            raise ValueError("Intentional error")

        with pytest.raises(ValueError):
            async with stream_with_context(failing_stream(), session_id="test") as stream:
                async for _ in stream:
                    pass

        # Context should be cleaned up
        assert get_current_session_id() is None


# =============================================================================
# Integration Tests
# =============================================================================

class TestPhase2Integration:
    """Integration tests for Phase 2 fixes working together."""

    @pytest.mark.asyncio
    async def test_traced_stream_with_ttl_registry(self, sample_span):
        """traced_stream works with TTL-based registry."""
        async def source_stream():
            for i in range(5):
                yield f"chunk-{i}"

        register_span(sample_span, ttl_seconds=60.0)
        try:
            chunks = []
            async for chunk in traced_stream(sample_span.span_id, source_stream()):
                chunks.append(chunk)

            assert chunks == [f"chunk-{i}" for i in range(5)]
            assert sample_span.attributes.get("stream.total_chunks") == 5
        finally:
            unregister_span(sample_span.span_id)

    @pytest.mark.asyncio
    async def test_long_running_stream_with_ttl_extension(self, sample_span):
        """Long streams can extend TTL to stay alive."""
        async def slow_stream():
            for i in range(5):
                await asyncio.sleep(0.05)  # Simulate slow chunks
                yield i

        # Short initial TTL
        register_span(sample_span, ttl_seconds=0.1)
        try:
            chunks = []
            async for chunk in traced_stream(sample_span.span_id, slow_stream()):
                # Extend TTL on each chunk
                extend_span_ttl(sample_span.span_id, 1.0)
                chunks.append(chunk)

            assert chunks == [0, 1, 2, 3, 4]
        finally:
            unregister_span(sample_span.span_id)

    def test_concurrent_streams_with_registry(self):
        """Multiple concurrent streams don't interfere."""
        patch_executors()

        results = {}
        errors = []

        def stream_worker(worker_id):
            try:
                span = Span(
                    trace_id=f"{worker_id:032x}",  # Valid 32-char hex
                    span_id=f"{worker_id:016x}",   # Valid 16-char hex
                    name=f"span-{worker_id}",
                )
                register_span(span, ttl_seconds=60.0)

                # Simulate some work
                time.sleep(0.01)

                # Verify span is still there
                retrieved = get_span_by_id(span.span_id)
                results[worker_id] = retrieved is span

                unregister_span(span.span_id)
            except Exception as e:
                errors.append((worker_id, str(e)))

        with ThreadPoolExecutor(max_workers=10) as executor:
            futures = [executor.submit(stream_worker, i) for i in range(50)]
            for f in futures:
                f.result()

        unpatch_executors()

        assert len(errors) == 0, f"Errors: {errors}"
        assert all(results.values()), "Some spans were not found"
        assert len(results) == 50


class TestBackwardCompatibility:
    """Ensure Phase 2 changes don't break existing behavior."""

    def test_existing_span_registry_api(self, sample_span):
        """Existing span registry API still works."""
        # Old pattern should still work
        register_span(sample_span)
        assert get_span_by_id(sample_span.span_id) is sample_span
        unregister_span(sample_span.span_id)
        assert get_span_by_id(sample_span.span_id) is None

    @pytest.mark.asyncio
    async def test_existing_traced_stream_api(self, sample_span):
        """Existing traced_stream API still works."""
        async def source():
            for i in range(3):
                yield i

        register_span(sample_span)
        try:
            results = [x async for x in traced_stream(sample_span.span_id, source())]
            assert results == [0, 1, 2]
        finally:
            unregister_span(sample_span.span_id)
