"""
Tests for Phase 3 fixes: Validation & Safety.

These tests verify:
- Gap 6 Fix: TraceContext validation (strict and lenient parsing)
- Gap 1 Fix: Pre-init executor detection and warnings
- Gap 8 Verification: Cancellation handling preserves cleanup
"""

import asyncio
import logging
import threading
import pytest
from concurrent.futures import ThreadPoolExecutor, Future

from risicare_core import (
    # Context managers
    session_context,
    async_session_context,
    agent_context,
    # Accessors
    get_current_session_id,
    get_current_agent_id,
    # TraceContext
    TraceContext,
    get_trace_context,
    restore_trace_context,
    # Patching
    patch_executors,
    unpatch_executors,
    is_executors_patched,
    get_pre_patch_executor_count,
    set_pre_init_executor_warning,
    # Types
    Span,
    generate_trace_id,
    generate_span_id,
)


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

@pytest.fixture(autouse=True)
def cleanup_patching():
    """Ensure clean state before and after each test."""
    # Unpatch before test
    unpatch_executors()
    yield
    # Unpatch after test
    unpatch_executors()
    # Re-enable warnings
    set_pre_init_executor_warning(True)


# =============================================================================
# Gap 6 Tests: TraceContext Validation
# =============================================================================

class TestTraceContextValidation:
    """Tests for TraceContext validation (Gap 6 fix)."""

    def test_valid_trace_context_creation(self):
        """Valid trace context is created successfully."""
        ctx = TraceContext(
            trace_id="a" * 32,
            span_id="b" * 16,
            session_id="my-session",
            agent_id="my-agent",
        )
        assert ctx.trace_id == "a" * 32
        assert ctx.span_id == "b" * 16
        assert ctx.session_id == "my-session"
        assert ctx.agent_id == "my-agent"

    def test_invalid_trace_id_too_short(self):
        """Invalid trace_id (too short) raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext(
                trace_id="abc",
                span_id="b" * 16,
            )
        assert "Invalid trace_id" in str(exc_info.value)
        assert "32-character" in str(exc_info.value)

    def test_invalid_trace_id_too_long(self):
        """Invalid trace_id (too long) raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext(
                trace_id="a" * 40,
                span_id="b" * 16,
            )
        assert "Invalid trace_id" in str(exc_info.value)

    def test_invalid_trace_id_non_hex(self):
        """Invalid trace_id (non-hex chars) raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext(
                trace_id="g" * 32,  # 'g' is not hex
                span_id="b" * 16,
            )
        assert "Invalid trace_id" in str(exc_info.value)

    def test_invalid_trace_id_uppercase(self):
        """Invalid trace_id (uppercase) raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext(
                trace_id="A" * 32,  # Uppercase not allowed
                span_id="b" * 16,
            )
        assert "Invalid trace_id" in str(exc_info.value)
        assert "lowercase" in str(exc_info.value)

    def test_invalid_span_id_too_short(self):
        """Invalid span_id (too short) raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext(
                trace_id="a" * 32,
                span_id="b" * 10,
            )
        assert "Invalid span_id" in str(exc_info.value)
        assert "16-character" in str(exc_info.value)

    def test_invalid_span_id_non_hex(self):
        """Invalid span_id (non-hex chars) raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext(
                trace_id="a" * 32,
                span_id="z" * 16,  # 'z' is not hex
            )
        assert "Invalid span_id" in str(exc_info.value)


class TestTraceContextFromDict:
    """Tests for TraceContext.from_dict() validation."""

    def test_from_dict_valid(self):
        """from_dict with valid data succeeds."""
        ctx = TraceContext.from_dict({
            "trace_id": "a" * 32,
            "span_id": "b" * 16,
            "session_id": "sess-123",
            "agent_id": "agent-456",
        })
        assert ctx.trace_id == "a" * 32
        assert ctx.span_id == "b" * 16
        assert ctx.session_id == "sess-123"
        assert ctx.agent_id == "agent-456"

    def test_from_dict_minimal(self):
        """from_dict with only required fields succeeds."""
        ctx = TraceContext.from_dict({
            "trace_id": "a" * 32,
            "span_id": "b" * 16,
        })
        assert ctx.session_id is None
        assert ctx.agent_id is None

    def test_from_dict_missing_trace_id(self):
        """from_dict with missing trace_id raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext.from_dict({
                "span_id": "b" * 16,
            })
        assert "Missing required field 'trace_id'" in str(exc_info.value)

    def test_from_dict_missing_span_id(self):
        """from_dict with missing span_id raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext.from_dict({
                "trace_id": "a" * 32,
            })
        assert "Missing required field 'span_id'" in str(exc_info.value)

    def test_from_dict_invalid_trace_id_type(self):
        """from_dict with wrong type raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext.from_dict({
                "trace_id": 12345,  # Should be string
                "span_id": "b" * 16,
            })
        assert "Invalid trace_id type" in str(exc_info.value)

    def test_from_dict_invalid_span_id_type(self):
        """from_dict with wrong span_id type raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext.from_dict({
                "trace_id": "a" * 32,
                "span_id": ["not", "a", "string"],
            })
        assert "Invalid span_id type" in str(exc_info.value)

    def test_from_dict_invalid_format(self):
        """from_dict with invalid format raises ValueError."""
        with pytest.raises(ValueError) as exc_info:
            TraceContext.from_dict({
                "trace_id": "not-valid-hex!!!",
                "span_id": "b" * 16,
            })
        assert "Invalid trace_id" in str(exc_info.value)


class TestTraceContextFromDictLenient:
    """Tests for TraceContext.from_dict_lenient() graceful handling."""

    def test_from_dict_lenient_valid(self):
        """from_dict_lenient with valid data returns TraceContext."""
        ctx = TraceContext.from_dict_lenient({
            "trace_id": "a" * 32,
            "span_id": "b" * 16,
        })
        assert ctx is not None
        assert ctx.trace_id == "a" * 32

    def test_from_dict_lenient_missing_trace_id(self):
        """from_dict_lenient with missing trace_id returns None."""
        ctx = TraceContext.from_dict_lenient({
            "span_id": "b" * 16,
        })
        assert ctx is None

    def test_from_dict_lenient_invalid_format(self):
        """from_dict_lenient with invalid format returns None."""
        ctx = TraceContext.from_dict_lenient({
            "trace_id": "invalid",
            "span_id": "also-invalid",
        })
        assert ctx is None

    def test_from_dict_lenient_wrong_type(self):
        """from_dict_lenient with wrong types returns None."""
        ctx = TraceContext.from_dict_lenient({
            "trace_id": 123,
            "span_id": 456,
        })
        assert ctx is None

    def test_from_dict_lenient_empty_dict(self):
        """from_dict_lenient with empty dict returns None."""
        ctx = TraceContext.from_dict_lenient({})
        assert ctx is None

    def test_from_dict_lenient_none_values(self):
        """from_dict_lenient with None values returns None."""
        ctx = TraceContext.from_dict_lenient({
            "trace_id": None,
            "span_id": None,
        })
        assert ctx is None


class TestTraceContextSerialization:
    """Tests for TraceContext serialization round-trip."""

    def test_to_dict_and_back(self):
        """to_dict and from_dict round-trip preserves data."""
        original = TraceContext(
            trace_id="a" * 32,
            span_id="b" * 16,
            session_id="my-session",
            agent_id="my-agent",
        )
        data = original.to_dict()
        restored = TraceContext.from_dict(data)

        assert restored.trace_id == original.trace_id
        assert restored.span_id == original.span_id
        assert restored.session_id == original.session_id
        assert restored.agent_id == original.agent_id

    def test_get_trace_context_generates_valid(self):
        """get_trace_context generates valid IDs when no context."""
        ctx = get_trace_context()
        # Should not raise - IDs are validated in TraceContext.__post_init__
        assert len(ctx.trace_id) == 32
        assert len(ctx.span_id) == 16
        # Verify they're valid hex
        int(ctx.trace_id, 16)
        int(ctx.span_id, 16)


# =============================================================================
# Gap 1 Tests: Pre-init Executor Detection
# =============================================================================

class TestPreInitExecutorDetection:
    """Tests for pre-init executor detection (Gap 1 fix)."""

    def test_executor_before_patch_is_tracked(self):
        """Executors created before patching are tracked."""
        # Create executor BEFORE patching
        executor = ThreadPoolExecutor(max_workers=1)

        # Now patch
        patch_executors()

        # Check that the pre-init executor was tracked
        assert get_pre_patch_executor_count() >= 1

        executor.shutdown(wait=False)

    def test_executor_after_patch_not_tracked(self):
        """Executors created after patching are not tracked as pre-init."""
        # Patch first
        patch_executors()

        # Get count before creating new executor
        count_before = get_pre_patch_executor_count()

        # Create executor AFTER patching
        executor = ThreadPoolExecutor(max_workers=1)

        # Count should be the same (new executor not tracked as pre-init)
        count_after = get_pre_patch_executor_count()
        assert count_after == count_before

        executor.shutdown(wait=False)

    def test_pre_init_executor_warning_logged(self, caplog):
        """Using pre-init executor logs a warning."""
        # Create executor BEFORE patching
        executor = ThreadPoolExecutor(max_workers=1)

        # Patch
        patch_executors()

        # Enable warning logging
        with caplog.at_level(logging.WARNING, logger="risicare.context.executors"):
            # Submit a task - should trigger warning
            future = executor.submit(lambda: 42)
            result = future.result()

        # Check warning was logged
        assert any("BEFORE patch_executors()" in msg for msg in caplog.messages)
        assert any("Context will NOT propagate" in msg for msg in caplog.messages)

        executor.shutdown(wait=False)

    def test_pre_init_executor_warning_only_once(self, caplog):
        """Warning for pre-init executor is logged only once per executor."""
        # Create executor BEFORE patching
        executor = ThreadPoolExecutor(max_workers=1)

        # Patch
        patch_executors()

        with caplog.at_level(logging.WARNING, logger="risicare.context.executors"):
            # Submit multiple tasks
            for _ in range(5):
                future = executor.submit(lambda: 42)
                future.result()

        # Should only have 1 warning (not 5)
        warning_count = sum(
            1 for msg in caplog.messages if "BEFORE patch_executors()" in msg
        )
        assert warning_count == 1

        executor.shutdown(wait=False)

    def test_warning_can_be_disabled(self, caplog):
        """Pre-init executor warning can be disabled."""
        # Create executor BEFORE patching
        executor = ThreadPoolExecutor(max_workers=1)

        # Patch
        patch_executors()

        # Disable warning
        set_pre_init_executor_warning(False)

        with caplog.at_level(logging.WARNING, logger="risicare.context.executors"):
            future = executor.submit(lambda: 42)
            future.result()

        # No warning should be logged
        assert not any("BEFORE patch_executors()" in msg for msg in caplog.messages)

        executor.shutdown(wait=False)

    def test_post_patch_executor_no_warning(self, caplog):
        """Executors created after patching don't trigger warning."""
        # Patch first
        patch_executors()

        # Create executor AFTER patching
        executor = ThreadPoolExecutor(max_workers=1)

        with caplog.at_level(logging.WARNING, logger="risicare.context.executors"):
            future = executor.submit(lambda: 42)
            future.result()

        # No warning should be logged
        assert not any("BEFORE patch_executors()" in msg for msg in caplog.messages)

        executor.shutdown(wait=False)

    def test_context_propagates_despite_warning(self):
        """Context still propagates for pre-init executors (after patching)."""
        # Create executor BEFORE patching
        executor = ThreadPoolExecutor(max_workers=1)

        # Patch (this patches the submit method)
        patch_executors()
        set_pre_init_executor_warning(False)  # Suppress warning for this test

        result = {}
        def capture_context():
            result["session_id"] = get_current_session_id()

        with session_context("test-session"):
            future = executor.submit(capture_context)
            future.result()

        # Context DOES propagate because submit is patched at class level
        assert result["session_id"] == "test-session"

        executor.shutdown(wait=False)


# =============================================================================
# Gap 8 Tests: Cancellation Handling
# =============================================================================

class TestCancellationHandling:
    """Tests for proper cleanup on cancellation (Gap 8)."""

    @pytest.mark.asyncio
    async def test_async_session_cancellation_cleanup(self):
        """Async session context cleans up on cancellation."""
        async def long_running():
            async with async_session_context("cancelled-session"):
                await asyncio.sleep(10)  # Will be cancelled

        task = asyncio.create_task(long_running())
        await asyncio.sleep(0.01)  # Let it start
        task.cancel()

        try:
            await task
        except asyncio.CancelledError:
            pass

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

    @pytest.mark.asyncio
    async def test_nested_context_cancellation(self):
        """Nested contexts clean up properly on cancellation."""
        async def nested_contexts():
            async with async_session_context("outer-session"):
                with agent_context("inner-agent"):
                    await asyncio.sleep(10)  # Will be cancelled

        task = asyncio.create_task(nested_contexts())
        await asyncio.sleep(0.01)
        task.cancel()

        try:
            await task
        except asyncio.CancelledError:
            pass

        # Both contexts should be cleaned up
        assert get_current_session_id() is None
        assert get_current_agent_id() is None

    @pytest.mark.asyncio
    async def test_exception_in_context_cleanup(self):
        """Exception in context block still cleans up."""
        with pytest.raises(ValueError):
            async with async_session_context("error-session"):
                raise ValueError("Intentional error")

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

    @pytest.mark.asyncio
    async def test_multiple_task_cancellation(self):
        """Multiple tasks can be cancelled without leaking context."""
        async def worker(session_id: str):
            async with async_session_context(session_id):
                await asyncio.sleep(10)

        tasks = [
            asyncio.create_task(worker(f"session-{i}"))
            for i in range(5)
        ]

        await asyncio.sleep(0.01)

        for task in tasks:
            task.cancel()

        await asyncio.gather(*tasks, return_exceptions=True)

        # No context should leak
        assert get_current_session_id() is None


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

class TestPhase3Integration:
    """Integration tests for Phase 3 fixes working together."""

    def test_trace_context_with_restore(self):
        """TraceContext validation works with restore_trace_context."""
        with session_context("original-session"):
            with agent_context("original-agent"):
                # Capture context
                ctx = get_trace_context()

                # Validate context fields
                assert len(ctx.trace_id) == 32
                assert len(ctx.span_id) == 16
                assert ctx.session_id == "original-session"
                assert ctx.agent_id == "original-agent"

        # Restore in different thread
        patch_executors()
        executor = ThreadPoolExecutor(max_workers=1)

        result = {}
        def restore_and_check(trace_ctx: TraceContext):
            with restore_trace_context(trace_ctx):
                result["session"] = get_current_session_id()
                result["agent"] = get_current_agent_id()

        future = executor.submit(restore_and_check, ctx)
        future.result()

        assert result["session"] == "original-session"
        assert result["agent"] == "original-agent"

        executor.shutdown(wait=False)

    def test_invalid_context_rejected_at_boundary(self):
        """Invalid context data is rejected at deserialization boundary."""
        # Simulate receiving bad data from external source
        bad_data = {
            "trace_id": "not-valid-hex-at-all!!!!",
            "span_id": "also-bad",
        }

        # Lenient parsing returns None
        ctx = TraceContext.from_dict_lenient(bad_data)
        assert ctx is None

        # Strict parsing raises
        with pytest.raises(ValueError):
            TraceContext.from_dict(bad_data)

    @pytest.mark.asyncio
    async def test_validation_and_cancellation_combined(self):
        """TraceContext validation works with async cancellation."""
        captured_ctx = None

        async def capture_and_wait():
            nonlocal captured_ctx
            async with async_session_context("capture-session"):
                captured_ctx = get_trace_context()
                await asyncio.sleep(10)  # Will be cancelled

        task = asyncio.create_task(capture_and_wait())
        await asyncio.sleep(0.01)
        task.cancel()

        try:
            await task
        except asyncio.CancelledError:
            pass

        # Context should have been captured before cancellation
        assert captured_ctx is not None
        assert captured_ctx.session_id == "capture-session"

        # And cleaned up after
        assert get_current_session_id() is None
