"""Integration tests for Mode A with Gemini: Flowlines bootstraps everything."""

from __future__ import annotations

from typing import TYPE_CHECKING

import pytest
from opentelemetry import trace

import flowlines

from .conftest import (
    OTLPCaptureServer,
    _skip_unless_gemini,
    make_gemini_call,
    make_gemini_chat_call,
    make_gemini_chat_stream_call,
)

if TYPE_CHECKING:
    from collections.abc import Callable


@pytest.fixture(autouse=True)
def _require_gemini() -> None:
    """Skip unless Gemini API key and package are available."""
    _skip_unless_gemini()


def test_spans_exported_to_flowlines_server(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode A: LLM spans arrive at the capture server with expected attributes."""
    server = capture_server()

    flowlines.init(api_key="test-key", ingest_endpoint=server.endpoint)
    try:
        make_gemini_call(user_id="user-42", session_id="convo-abc")
    finally:
        flowlines.shutdown()

    spans = server.wait_for_spans(min_count=1)
    assert len(spans) >= 1, f"Expected ≥1 span, got {len(spans)}"

    # Verify LLM-specific attributes are present
    span = spans[0]
    attr_keys = {a["key"] for a in span.get("attributes", [])}
    gen_ai_keys = {k for k in attr_keys if k.startswith("gen_ai.")}
    assert gen_ai_keys, f"No gen_ai.* attributes found. Keys: {attr_keys}"

    # Verify model name
    attr_map = {a["key"]: a["value"] for a in span.get("attributes", [])}
    model_attr = attr_map.get("gen_ai.request.model", {})
    model_value = model_attr.get("string_value") or model_attr.get("stringValue", "")
    assert "gemini" in model_value, f"Unexpected model: {model_value}"

    # Verify Flowlines context 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-42", f"Unexpected user_id: {user_id_value}"

    convo_attr = attr_map.get("flowlines.session_id", {})
    convo_value = convo_attr.get("string_value") or convo_attr.get("stringValue", "")
    assert convo_value == "convo-abc", f"Unexpected session_id: {convo_value}"


def test_api_key_sent_in_header(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode A: The x-flowlines-api-key header is present in OTLP export requests."""
    server = capture_server()

    flowlines.init(api_key="my-secret-key", ingest_endpoint=server.endpoint)
    try:
        make_gemini_call(user_id="user-42", session_id="convo-abc")
    finally:
        flowlines.shutdown()

    server.wait_for_spans(min_count=1)
    records = server.get_records()
    assert len(records) >= 1

    headers = records[0]["headers"]
    assert headers.get("x-flowlines-api-key") == "my-secret-key"


def test_non_llm_spans_are_filtered_out(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode A: non-LLM spans are filtered out and never reach the capture server."""
    server = capture_server()

    flowlines.init(api_key="test-key", ingest_endpoint=server.endpoint)
    try:
        # Create a non-LLM span via the global tracer provider
        tracer = trace.get_tracer("test-non-llm")
        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 a real LLM span via Gemini
        make_gemini_call(user_id="user-1", session_id="convo-1")
    finally:
        flowlines.shutdown()

    spans = server.wait_for_spans(min_count=1)
    assert len(spans) >= 1, f"Expected ≥1 span, got {len(spans)}"

    # Every span that reached the server must have gen_ai.* or ai.* attributes
    for span in spans:
        attr_keys = {a["key"] for a in span.get("attributes", [])}
        llm_keys = {
            k for k in attr_keys if k.startswith("gen_ai.") or k.startswith("ai.")
        }
        assert llm_keys, (
            f"Span '{span.get('name')}' has no gen_ai.*/ai.*"
            f" attributes. Keys: {attr_keys}"
        )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _get_attr(span: dict, key: str) -> str:
    """Extract a string attribute value from a captured span."""
    attr_map = {a["key"]: a["value"] for a in span.get("attributes", [])}
    val = attr_map.get(key, {})
    return val.get("string_value") or val.get("stringValue", "")


# ---------------------------------------------------------------------------
# Tests: user messages (prompts) and assistant messages (completions)
# ---------------------------------------------------------------------------


def test_spans_contain_user_and_assistant_messages(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode A (string contents): spans carry both prompt and completion attributes."""
    server = capture_server()

    flowlines.init(api_key="test-key", ingest_endpoint=server.endpoint)
    try:
        make_gemini_call(user_id="u1", session_id="s1")
    finally:
        flowlines.shutdown()

    spans = server.wait_for_spans(min_count=1)
    assert len(spans) >= 1
    span = spans[0]

    # User message (prompt)
    prompt_content = _get_attr(span, "gen_ai.prompt.0.content")
    assert prompt_content, (
        f"No gen_ai.prompt.0.content attribute found. Keys: "
        f"{sorted(a['key'] for a in span.get('attributes', []))}"
    )
    prompt_role = _get_attr(span, "gen_ai.prompt.0.role")
    assert prompt_role == "user", f"Expected role 'user', got '{prompt_role}'"

    # Assistant message (completion)
    completion_content = _get_attr(span, "gen_ai.completion.0.content")
    assert completion_content, (
        f"No gen_ai.completion.0.content attribute found. Keys: "
        f"{sorted(a['key'] for a in span.get('attributes', []))}"
    )
    completion_role = _get_attr(span, "gen_ai.completion.0.role")
    assert completion_role == "assistant", (
        f"Expected role 'assistant', got '{completion_role}'"
    )


def test_chat_api_spans_contain_user_and_assistant_messages(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode A (Chat API, list contents): spans carry both prompt and completion.

    The Chat API always passes ``contents`` as a list of ``Content`` objects,
    which exercises a different code path in the instrumentor than a plain
    string. This test verifies that the LLM_PROMPTS monkey-patch is effective.
    """
    server = capture_server()

    flowlines.init(api_key="test-key", ingest_endpoint=server.endpoint)
    try:
        make_gemini_chat_call(user_id="u1", session_id="s1")
    finally:
        flowlines.shutdown()

    spans = server.wait_for_spans(min_count=1)
    assert len(spans) >= 1
    span = spans[0]

    # User message (prompt)
    prompt_content = _get_attr(span, "gen_ai.prompt.0.content")
    assert prompt_content, (
        f"No gen_ai.prompt.0.content attribute found. Keys: "
        f"{sorted(a['key'] for a in span.get('attributes', []))}"
    )
    prompt_role = _get_attr(span, "gen_ai.prompt.0.role")
    assert prompt_role == "user", f"Expected role 'user', got '{prompt_role}'"

    # Assistant message (completion)
    completion_content = _get_attr(span, "gen_ai.completion.0.content")
    assert completion_content, (
        f"No gen_ai.completion.0.content attribute found. Keys: "
        f"{sorted(a['key'] for a in span.get('attributes', []))}"
    )
    completion_role = _get_attr(span, "gen_ai.completion.0.role")
    assert completion_role == "assistant", (
        f"Expected role 'assistant', got '{completion_role}'"
    )


def test_chat_stream_api_spans_contain_user_and_assistant_messages(
    capture_server: Callable[[], OTLPCaptureServer],
) -> None:
    """Mode A (Chat streaming API): spans carry both prompt and completion.

    ``chat.send_message_stream()`` exercises both the list-of-contents path
    *and* the streaming response path. This matches the real-world usage
    pattern that originally surfaced the missing user-messages bug.
    """
    server = capture_server()

    flowlines.init(api_key="test-key", ingest_endpoint=server.endpoint)
    try:
        make_gemini_chat_stream_call(user_id="u1", session_id="s1")
    finally:
        flowlines.shutdown()

    spans = server.wait_for_spans(min_count=1)
    assert len(spans) >= 1
    span = spans[0]

    # User message (prompt)
    prompt_content = _get_attr(span, "gen_ai.prompt.0.content")
    assert prompt_content, (
        f"No gen_ai.prompt.0.content attribute found. Keys: "
        f"{sorted(a['key'] for a in span.get('attributes', []))}"
    )
    prompt_role = _get_attr(span, "gen_ai.prompt.0.role")
    assert prompt_role == "user", f"Expected role 'user', got '{prompt_role}'"

    # Assistant message (completion)
    completion_content = _get_attr(span, "gen_ai.completion.0.content")
    assert completion_content, (
        f"No gen_ai.completion.0.content attribute found. Keys: "
        f"{sorted(a['key'] for a in span.get('attributes', []))}"
    )
    completion_role = _get_attr(span, "gen_ai.completion.0.role")
    assert completion_role == "assistant", (
        f"Expected role 'assistant', got '{completion_role}'"
    )
