"""
Tests for context propagation observability.

Verifies:
- Metrics tracking (counters, histograms)
- Context loss detection
- Logging integration
- Health reporting
- Prometheus export
"""

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

from risicare_core import (
    # Context managers
    session_context,
    async_session_context,
    agent_context,
    async_agent_context,
    phase_context,
    # Accessors
    get_current_session,
    get_current_agent,
    get_current_span,
    # Patching
    patch_executors,
    unpatch_executors,
)
from risicare_core.observability import (
    # Configuration
    enable_context_observability,
    disable_context_observability,
    is_observability_enabled,
    # Types
    ContextType,
    ContextOperation,
    # Metrics
    context_metrics,
    # Scope tracking
    enter_traced_scope,
    exit_traced_scope,
    is_inside_traced_scope,
    # Reporting
    get_context_health_report,
    get_prometheus_metrics,
)
from risicare_core.types.base import SemanticPhase


# =============================================================================
# Setup / Teardown
# =============================================================================

@pytest.fixture(autouse=True)
def reset_observability():
    """Reset observability state before and after each test."""
    context_metrics.reset_all()
    disable_context_observability()
    yield
    context_metrics.reset_all()
    disable_context_observability()


@pytest.fixture
def enable_obs():
    """Enable observability for a test."""
    enable_context_observability(
        log_level="DEBUG",
        log_context_enter_exit=True,
        track_timing=True,
    )
    return context_metrics


# =============================================================================
# Configuration Tests
# =============================================================================

class TestConfiguration:
    def test_enable_disable(self):
        """Observability can be enabled and disabled."""
        assert not is_observability_enabled()

        enable_context_observability()
        assert is_observability_enabled()

        disable_context_observability()
        assert not is_observability_enabled()

    def test_enable_with_options(self):
        """Observability accepts configuration options."""
        enable_context_observability(
            log_level="DEBUG",
            log_context_enter_exit=True,
            track_timing=True,
            sample_rate=0.5,
        )
        assert is_observability_enabled()


# =============================================================================
# Counter Tests
# =============================================================================

class TestCounters:
    def test_counter_basic_operations(self, enable_obs):
        """Counters track operations correctly."""
        # Use session context
        with session_context("test-session"):
            pass

        # Check operation counts
        ops = context_metrics.context_operations.get_all()

        # Should have enter and exit for session
        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "enter",
        }) == 1
        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "exit",
        }) == 1

    def test_counter_nested_contexts(self, enable_obs):
        """Counters track nested context operations."""
        with session_context("session-1"):
            with agent_context("agent-1"):
                with phase_context(SemanticPhase.THINK):
                    pass

        # Each context type should have enter and exit
        for ctx_type in ["session", "agent", "phase"]:
            assert context_metrics.context_operations.get({
                "context_type": ctx_type,
                "operation": "enter",
            }) == 1
            assert context_metrics.context_operations.get({
                "context_type": ctx_type,
                "operation": "exit",
            }) == 1

    def test_counter_multiple_operations(self, enable_obs):
        """Counters accumulate across multiple operations."""
        for i in range(5):
            with session_context(f"session-{i}"):
                pass

        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "enter",
        }) == 5
        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "exit",
        }) == 5

    def test_counter_reset(self, enable_obs):
        """Counters can be reset."""
        with session_context("test"):
            pass

        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "enter",
        }) == 1

        context_metrics.reset_all()

        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "enter",
        }) == 0


# =============================================================================
# Context Loss Detection Tests
# =============================================================================

class TestContextLossDetection:
    def test_scope_tracking(self):
        """Traced scope tracking works correctly."""
        assert not is_inside_traced_scope()

        enter_traced_scope()
        assert is_inside_traced_scope()

        exit_traced_scope()
        assert not is_inside_traced_scope()

    def test_scope_nesting(self):
        """Traced scope supports nesting."""
        assert not is_inside_traced_scope()

        enter_traced_scope()
        enter_traced_scope()
        assert is_inside_traced_scope()

        exit_traced_scope()
        assert is_inside_traced_scope()  # Still nested

        exit_traced_scope()
        assert not is_inside_traced_scope()

    def test_context_loss_recorded_when_expected(self, enable_obs):
        """Context loss is recorded when inside traced scope."""
        # Enter traced scope manually (simulating being in context)
        enter_traced_scope()

        # Try to get session when none exists
        session = get_current_session()

        exit_traced_scope()

        # Should have recorded a loss
        assert session is None
        assert context_metrics.context_loss.get({
            "context_type": "session",
            "expected": "true",
        }) == 1

    def test_no_loss_when_not_expected(self, enable_obs):
        """Context None is not recorded as loss when not in traced scope."""
        # Get session when not in any traced scope
        session = get_current_session()

        # Should not record loss (we're not expecting context)
        assert session is None
        assert context_metrics.loss_tracker.get_total_losses() == 0

    def test_loss_tracker_records_location(self, enable_obs):
        """Loss tracker records where loss occurred."""
        enter_traced_scope()
        get_current_agent()
        exit_traced_scope()

        events = context_metrics.loss_tracker.get_recent_events(1)
        assert len(events) == 1
        assert events[0].context_type == ContextType.AGENT
        assert "test_observability" in events[0].location

    def test_loss_tracker_top_locations(self, enable_obs):
        """Loss tracker identifies top loss locations."""
        enter_traced_scope()

        # Cause losses at same location multiple times
        for _ in range(5):
            get_current_session()

        exit_traced_scope()

        top = context_metrics.loss_tracker.get_top_locations(1)
        assert len(top) == 1
        assert top[0][1] == 5  # Count of 5


# =============================================================================
# Timing Tests
# =============================================================================

class TestTiming:
    def test_context_duration_recorded(self, enable_obs):
        """Context manager duration is recorded."""
        import time

        with session_context("timed-session"):
            time.sleep(0.01)  # 10ms

        assert context_metrics.context_duration.get_count() > 0
        p50 = context_metrics.context_duration.get_percentile(50)
        assert p50 >= 0.01  # At least 10ms

    def test_histogram_percentiles(self, enable_obs):
        """Histogram calculates percentiles correctly."""
        # Manually add observations
        for v in [0.01, 0.02, 0.03, 0.04, 0.05]:
            context_metrics.context_duration.observe(v)

        p50 = context_metrics.context_duration.get_percentile(50)
        assert 0.02 <= p50 <= 0.04  # Middle value


# =============================================================================
# Async Tests
# =============================================================================

class TestAsyncObservability:
    @pytest.mark.asyncio
    async def test_async_session_context_tracked(self, enable_obs):
        """Async session context is tracked."""
        async with async_session_context("async-session"):
            await asyncio.sleep(0.001)

        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "enter",
        }) == 1
        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "exit",
        }) == 1

    @pytest.mark.asyncio
    async def test_async_agent_context_tracked(self, enable_obs):
        """Async agent context is tracked."""
        async with async_agent_context("async-agent"):
            await asyncio.sleep(0.001)

        assert context_metrics.context_operations.get({
            "context_type": "agent",
            "operation": "enter",
        }) == 1

    @pytest.mark.asyncio
    async def test_async_cancellation_cleanup(self, enable_obs):
        """Async context cleanup on cancellation is tracked."""
        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)
        task.cancel()

        try:
            await task
        except asyncio.CancelledError:
            pass

        # Exit should still be recorded even though cancelled
        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "exit",
        }) >= 1


# =============================================================================
# Threading Tests
# =============================================================================

class TestThreadingObservability:
    def test_thread_safe_counters(self, enable_obs):
        """Counters are thread-safe."""
        patch_executors()

        def worker(i):
            with session_context(f"session-{i}"):
                pass

        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = [executor.submit(worker, i) for i in range(20)]
            for f in futures:
                f.result()

        unpatch_executors()

        # All 20 sessions should be tracked
        enter_count = context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "enter",
        })
        exit_count = context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "exit",
        })

        assert enter_count == 20
        assert exit_count == 20

    def test_loss_tracker_thread_safe(self, enable_obs):
        """Loss tracker is thread-safe."""
        def cause_loss(i):
            enter_traced_scope()
            get_current_session()
            exit_traced_scope()

        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = [executor.submit(cause_loss, i) for i in range(20)]
            for f in futures:
                f.result()

        # All 20 losses should be recorded
        assert context_metrics.loss_tracker.get_total_losses() == 20


# =============================================================================
# Health Report Tests
# =============================================================================

class TestHealthReport:
    def test_health_report_structure(self, enable_obs):
        """Health report has expected structure."""
        with session_context("test"):
            pass

        report = get_context_health_report()

        assert "enabled" in report
        assert "total_operations" in report
        assert "total_losses" in report
        assert "loss_rate" in report
        assert "timing" in report
        assert "top_loss_locations" in report
        assert "recent_loss_events" in report

    def test_health_report_loss_rate(self, enable_obs):
        """Health report calculates loss rate correctly."""
        # Create some operations
        for _ in range(10):
            with session_context("test"):
                pass

        # Create some losses
        for _ in range(2):
            enter_traced_scope()
            get_current_session()  # Loss
            exit_traced_scope()

        report = get_context_health_report()

        # Loss rate should be 2 / (20 + 2) = approx 0.083
        # (20 enter/exit ops + 2 loss recordings)
        assert report["total_losses"] == 2
        assert report["loss_rate"] > 0


# =============================================================================
# Prometheus Export Tests
# =============================================================================

class TestPrometheusExport:
    def test_prometheus_export_without_client(self):
        """Prometheus export returns None without prometheus_client."""
        # This may or may not have prometheus_client installed
        result = get_prometheus_metrics()
        # Just verify it doesn't crash
        assert result is None or isinstance(result, str)

    def test_prometheus_export_with_data(self, enable_obs):
        """Prometheus export includes recorded metrics."""
        with session_context("test"):
            pass

        result = get_prometheus_metrics()

        # If prometheus_client is installed, check output
        if result is not None:
            assert "risicare_context_operations_total" in result
            assert "session" in result


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

class TestIntegration:
    def test_full_observability_flow(self, enable_obs, caplog):
        """Test complete observability flow."""
        with caplog.at_level(logging.DEBUG, logger="risicare.context"):
            # Normal operation
            with session_context("integration-session", user_id="test-user"):
                with agent_context("integration-agent", agent_role="worker"):
                    # This should not cause loss (we're in context)
                    session = get_current_session()
                    agent = get_current_agent()

            # Verify we got context
            assert session is not None
            assert agent is not None

            # Cause a loss
            enter_traced_scope()
            lost_span = get_current_span()
            exit_traced_scope()

        # Verify metrics
        report = get_context_health_report()
        assert report["total_operations"] > 0
        assert report["total_losses"] == 1

    def test_observability_disabled_no_overhead(self):
        """When disabled, observability adds minimal overhead."""
        disable_context_observability()

        # Operations should still work
        with session_context("no-obs-session"):
            with agent_context("no-obs-agent"):
                session = get_current_session()

        assert session is not None

        # But no metrics recorded
        assert context_metrics.context_operations.get({
            "context_type": "session",
            "operation": "enter",
        }) == 0
