"""E2E tests for trace flow: SDK → Collector → GCP Cloud Trace.

Tests verify that distributed traces recorded via SDK successfully propagate
through the OTel Collector and appear in GCP Cloud Trace with proper span
hierarchy and timing information.
"""

import time
import pytest
from mca_sdk import MCAClient


@pytest.mark.integration
@pytest.mark.error_scenario  # Collector logs are expected
class TestTracesE2E:
    """End-to-end tests for distributed tracing."""

    def test_simple_trace_propagates_to_gcp(
        self,
        docker_compose_environment,
        gcp_trace_client,
        gcp_project_id,
        wait_for_telemetry,
        cleanup_collector_exports,
    ):
        """Test AC3: Single span trace flows to GCP Cloud Trace.

        Given: SDK creates trace for prediction
        When: Trace propagates through collector
        Then: Complete span is visible in GCP Cloud Trace
        """
        # Arrange
        client = MCAClient(
            service_name="e2e-trace-test",
            model_id="e2e-trace-mdl-001",
            team_name="e2e-test-team",
            model_category="internal",
            model_type="regression",
            otel_endpoint=docker_compose_environment["collector_otlp"],
            otel_protocol="http",
        )

        # Act: Create traced function
        trace_id = None

        @client.predict()
        def make_prediction(features: dict) -> dict:
            """Traced prediction function."""
            time.sleep(0.05)  # Simulate inference
            return {"prediction": 0.75, "confidence": 0.9}

        # Execute traced function and extract trace ID
        tracer = client._tracer
        with tracer.start_as_current_span("test_span") as span:
            trace_id = format(span.get_span_context().trace_id, "032x")  # Extract trace ID as hex
            result = make_prediction({"feature1": 1.0, "feature2": 2.0})

        client.shutdown()
        wait_for_telemetry(10, "trace propagation to file exporter")

        # Assert: Verify trace exists in GCP Cloud Trace
        assert trace_id is not None, "Failed to extract trace_id from span context"

        trace = gcp_trace_client.get_trace(gcp_project_id, trace_id)
        assert trace is not None, f"Trace {trace_id} not found in Cloud Trace exports"
        assert len(trace["spans"]) >= 1, f"Expected at least 1 span, found {len(trace['spans'])}"

        # Verify span has expected name
        span_names = [s["name"] for s in trace["spans"]]
        assert (
            "test_span" in span_names or "make_prediction" in span_names
        ), f"Expected span name not found in: {span_names}"

    def test_nested_spans_hierarchy_preserved(
        self,
        docker_compose_environment,
        gcp_trace_client,
        gcp_project_id,
        wait_for_telemetry,
        cleanup_collector_exports,
    ):
        """Test AC3: Nested span hierarchy is preserved in Cloud Trace.

        Given: SDK creates trace with nested spans (preprocessing → inference → postprocessing)
        When: Trace propagates to GCP
        Then: Span hierarchy and parent-child relationships are correct
        """
        # Arrange
        client = MCAClient(
            service_name="e2e-nested-trace",
            model_id="e2e-trace-mdl-002",
            team_name="e2e-test-team",
            model_category="internal",
            model_type="classification",
            otel_endpoint=docker_compose_environment["collector_otlp"],
            otel_protocol="http",
        )

        # Act: Create nested spans
        tracer = client._tracer
        trace_id = None

        with tracer.start_as_current_span("prediction") as parent_span:
            trace_id = format(parent_span.get_span_context().trace_id, "032x")

            with tracer.start_as_current_span("preprocessing"):
                time.sleep(0.02)

            with tracer.start_as_current_span("inference"):
                time.sleep(0.05)

            with tracer.start_as_current_span("postprocessing"):
                time.sleep(0.01)

        client.shutdown()
        wait_for_telemetry(10, "nested trace propagation")

        # Assert: Verify span hierarchy
        assert trace_id is not None, "Failed to extract trace_id"

        trace = gcp_trace_client.get_trace(gcp_project_id, trace_id)
        assert trace is not None, f"Nested trace {trace_id} not found"
        assert (
            len(trace["spans"]) == 4
        ), f"Expected 4 spans (1 parent + 3 children), found {len(trace['spans'])}"

        # Verify span names
        span_names = [span["name"] for span in trace["spans"]]
        assert "prediction" in span_names, f"Parent span 'prediction' not found in {span_names}"
        assert "preprocessing" in span_names, "Child span 'preprocessing' not found"
        assert "inference" in span_names, "Child span 'inference' not found"
        assert "postprocessing" in span_names, "Child span 'postprocessing' not found"

        # Verify parent-child relationships
        parent_span = next(s for s in trace["spans"] if s["name"] == "prediction")
        child_spans = [s for s in trace["spans"] if s["name"] != "prediction"]

        for child in child_spans:
            # Child should have parent_span_id matching parent's span_id
            assert child.get("parent_span_id") == parent_span["span_id"], (
                f"Child span '{child['name']}' parent_span_id mismatch. "
                f"Expected {parent_span['span_id']}, got {child.get('parent_span_id')}"
            )

    def test_span_timing_accuracy(
        self,
        docker_compose_environment,
        gcp_trace_client,
        gcp_project_id,
        wait_for_telemetry,
        cleanup_collector_exports,
    ):
        """Test that span timing information is reasonably accurate.

        Given: SDK creates spans with known durations
        When: Trace is retrieved from Cloud Trace
        Then: Span durations are within acceptable tolerance (±100ms for CI stability)
        """
        # Arrange
        client = MCAClient(
            service_name="e2e-timing-test",
            model_id="e2e-trace-mdl-003",
            team_name="e2e-test-team",
            model_category="internal",
            model_type="regression",
            otel_endpoint=docker_compose_environment["collector_otlp"],
            otel_protocol="http",
        )

        # Act: Create span with known duration
        tracer = client._tracer
        expected_duration = 0.1  # 100ms
        trace_id = None

        with tracer.start_as_current_span("timed_operation") as span:
            trace_id = format(span.get_span_context().trace_id, "032x")
            start = time.time()
            time.sleep(expected_duration)
            actual_duration = time.time() - start

        client.shutdown()
        wait_for_telemetry(10, "timing data export")

        # Assert: Verify timing accuracy
        assert trace_id is not None, "Failed to extract trace_id"

        trace = gcp_trace_client.get_trace(gcp_project_id, trace_id)
        assert trace is not None, f"Timing trace {trace_id} not found"
        assert len(trace["spans"]) > 0, "No spans found in trace"

        span = trace["spans"][0]
        recorded_duration = span.get("duration_seconds", 0)

        # Relaxed tolerance: 100ms to account for CI thread scheduling
        # time.sleep() is imprecise, especially in virtualized/CI environments
        tolerance = 0.1  # 100ms
        lower_bound = actual_duration - tolerance
        upper_bound = actual_duration + tolerance

        assert lower_bound <= recorded_duration <= upper_bound, (
            f"Timing mismatch: recorded={recorded_duration:.3f}s, "
            f"actual={actual_duration:.3f}s, tolerance=±{tolerance}s. "
            f"Expected range: [{lower_bound:.3f}, {upper_bound:.3f}]"
        )

    def test_trace_attributes_preserved(
        self,
        docker_compose_environment,
        gcp_trace_client,
        gcp_project_id,
        wait_for_telemetry,
        cleanup_collector_exports,
    ):
        """Test that custom trace attributes are preserved.

        Given: Spans have custom attributes of various types
        When: Trace is retrieved from Cloud Trace
        Then: All attributes are present with correct types and values
        """
        # Arrange
        client = MCAClient(
            service_name="e2e-trace-attrs",
            model_id="e2e-trace-mdl-004",
            team_name="e2e-test-team",
            model_category="internal",
            model_type="classification",
            otel_endpoint=docker_compose_environment["collector_otlp"],
            otel_protocol="http",
        )

        # Act: Create span with custom attributes
        tracer = client._tracer
        test_attributes = {
            "custom.attr1": "value1",
            "custom.attr2": 42,
            "custom.attr3": 3.14,
            "custom.bool": True,
        }
        trace_id = None

        with tracer.start_as_current_span("attributed_span") as span:
            trace_id = format(span.get_span_context().trace_id, "032x")
            for key, value in test_attributes.items():
                span.set_attribute(key, value)

        client.shutdown()
        wait_for_telemetry(10, "trace attributes export")

        # Assert
        assert trace_id is not None, "Failed to extract trace_id"

        trace = gcp_trace_client.get_trace(gcp_project_id, trace_id)
        assert trace is not None, f"Attribute trace {trace_id} not found"
        assert len(trace["spans"]) > 0, "No spans in trace"

        span_attrs = trace["spans"][0].get("attributes", {})

        # Verify each attribute
        for key, expected_value in test_attributes.items():
            assert (
                key in span_attrs
            ), f"Attribute '{key}' not found. Available: {list(span_attrs.keys())}"

            actual_value = span_attrs[key]

            # Type-specific assertions
            if isinstance(expected_value, bool):
                assert (
                    actual_value == expected_value
                ), f"Boolean attribute '{key}' mismatch: {actual_value} != {expected_value}"
            elif isinstance(expected_value, int):
                assert (
                    actual_value == expected_value
                ), f"Integer attribute '{key}' mismatch: {actual_value} != {expected_value}"
            elif isinstance(expected_value, float):
                assert (
                    abs(actual_value - expected_value) < 0.01
                ), f"Float attribute '{key}' mismatch: {actual_value} != {expected_value}"
            else:
                assert (
                    actual_value == expected_value
                ), f"String attribute '{key}' mismatch: {actual_value} != {expected_value}"
