"""Integration tests for Mode B1: Flowlines with existing OTEL setup."""

from __future__ import annotations

from typing import TYPE_CHECKING

import pytest
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

import flowlines

from .conftest import OTLPCaptureServer, _skip_unless_openai, make_openai_call

if TYPE_CHECKING:
    from collections.abc import Callable


@pytest.fixture(autouse=True)
def _require_openai() -> None:
    _skip_unless_openai()


def test_all_exporters_receive_spans(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode B1: Flowlines, Sentry, and Traceloop servers all receive LLM spans."""
    flowlines_server = capture_server()
    sentry_server = capture_server()
    traceloop_server = capture_server()

    # Set up user-owned TracerProvider with two exporters
    provider = TracerProvider()

    sentry_processor = BatchSpanProcessor(
        OTLPSpanExporter(endpoint=f"{sentry_server.endpoint}/v1/traces")
    )
    provider.add_span_processor(sentry_processor)

    traceloop_processor = BatchSpanProcessor(
        OTLPSpanExporter(endpoint=f"{traceloop_server.endpoint}/v1/traces")
    )
    provider.add_span_processor(traceloop_processor)

    trace.set_tracer_provider(provider)

    # Initialize Flowlines in Mode B1
    flowlines.init(
        api_key="test-key",
        ingest_endpoint=flowlines_server.endpoint,
        has_external_otel=True,
    )
    flowlines_processor = flowlines.create_span_processor()
    provider.add_span_processor(flowlines_processor)

    for instrumentor in flowlines.get_instrumentors():
        instrumentor.instrument(tracer_provider=provider)

    try:
        make_openai_call(user_id="user-b1", session_id="convo-b1")
    finally:
        sentry_processor.shutdown()  # type: ignore[no-untyped-call]
        traceloop_processor.shutdown()  # type: ignore[no-untyped-call]
        flowlines.shutdown()

    # All three servers should receive spans
    for name, server in [
        ("flowlines", flowlines_server),
        ("sentry", sentry_server),
        ("traceloop", traceloop_server),
    ]:
        spans = server.wait_for_spans(min_count=1)
        assert len(spans) >= 1, f"{name} server: expected ≥1 span, got {len(spans)}"

    # Verify Flowlines context attributes on the Flowlines server spans
    fl_spans = flowlines_server.get_spans()
    attr_map = {a["key"]: a["value"] for a in fl_spans[0].get("attributes", [])}
    user_id_attr = attr_map.get("flowlines.user_id", {})
    user_id_value = user_id_attr.get("string_value") or user_id_attr.get(
        "stringValue", ""
    )
    assert user_id_value == "user-b1", f"Unexpected user_id: {user_id_value}"


def test_flowlines_only_exports_llm_spans(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode B1: Flowlines filters to LLM spans only; Sentry gets all spans."""
    flowlines_server = capture_server()
    sentry_server = capture_server()

    provider = TracerProvider()

    sentry_processor = BatchSpanProcessor(
        OTLPSpanExporter(endpoint=f"{sentry_server.endpoint}/v1/traces")
    )
    provider.add_span_processor(sentry_processor)

    trace.set_tracer_provider(provider)

    flowlines.init(
        api_key="test-key",
        ingest_endpoint=flowlines_server.endpoint,
        has_external_otel=True,
    )
    flowlines_processor = flowlines.create_span_processor()
    provider.add_span_processor(flowlines_processor)

    for instrumentor in flowlines.get_instrumentors():
        instrumentor.instrument(tracer_provider=provider)

    try:
        # Create a manual non-LLM span
        tracer = trace.get_tracer("test-tracer")
        with tracer.start_as_current_span("http-request") as span:
            span.set_attribute("http.method", "GET")
            span.set_attribute("http.url", "https://example.com")

        # Create an LLM span via OpenAI
        make_openai_call(user_id="user-b1", session_id="convo-b1")
    finally:
        sentry_processor.shutdown()  # type: ignore[no-untyped-call]
        flowlines.shutdown()

    # Sentry should get both the manual span and LLM spans
    sentry_spans = sentry_server.wait_for_spans(min_count=2)
    assert len(sentry_spans) >= 2, (
        f"Sentry: expected ≥2 spans (manual + LLM), got {len(sentry_spans)}"
    )

    # Flowlines should only get LLM spans (filtered by gen_ai.*/ai.* attributes)
    flowlines_spans = flowlines_server.wait_for_spans(min_count=1)
    for span in flowlines_spans:
        attr_keys = {a["key"] for a in span.get("attributes", [])}
        has_llm_attr = any(
            k.startswith("gen_ai.") or k.startswith("ai.") for k in attr_keys
        )
        assert has_llm_attr, (
            f"Flowlines received a non-LLM span. Attributes: {attr_keys}"
        )


def test_api_key_only_sent_to_flowlines(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode B1: The x-flowlines-api-key header is only sent to the Flowlines server."""
    flowlines_server = capture_server()
    sentry_server = capture_server()
    traceloop_server = capture_server()

    # Set up user-owned TracerProvider with two exporters
    provider = TracerProvider()

    sentry_processor = BatchSpanProcessor(
        OTLPSpanExporter(endpoint=f"{sentry_server.endpoint}/v1/traces")
    )
    provider.add_span_processor(sentry_processor)

    traceloop_processor = BatchSpanProcessor(
        OTLPSpanExporter(endpoint=f"{traceloop_server.endpoint}/v1/traces")
    )
    provider.add_span_processor(traceloop_processor)

    trace.set_tracer_provider(provider)

    # Initialize Flowlines in Mode B1
    flowlines.init(
        api_key="my-secret-key",
        ingest_endpoint=flowlines_server.endpoint,
        has_external_otel=True,
    )
    flowlines_processor = flowlines.create_span_processor()
    provider.add_span_processor(flowlines_processor)

    for instrumentor in flowlines.get_instrumentors():
        instrumentor.instrument(tracer_provider=provider)

    try:
        make_openai_call(user_id="user-b1", session_id="convo-b1")
    finally:
        sentry_processor.shutdown()  # type: ignore[no-untyped-call]
        traceloop_processor.shutdown()  # type: ignore[no-untyped-call]
        flowlines.shutdown()

    # Flowlines server must receive the API key header
    flowlines_server.wait_for_spans(min_count=1)
    fl_records = flowlines_server.get_records()
    assert len(fl_records) >= 1
    for record in fl_records:
        assert record["headers"].get("x-flowlines-api-key") == "my-secret-key", (
            "Flowlines server did not receive the expected x-flowlines-api-key header"
        )

    # Sentry and Traceloop servers must NOT receive the API key header
    for name, server in [("sentry", sentry_server), ("traceloop", traceloop_server)]:
        server.wait_for_spans(min_count=1)
        for record in server.get_records():
            assert "x-flowlines-api-key" not in record["headers"], (
                f"{name} server received x-flowlines-api-key header — credential leak!"
            )
