"""Harness coverage for the gen-AI span selection `logging/test_otel_trace_e2e` relies on.

This exercises the selection helper itself against Jaeger-shaped payloads, so it
runs without a proxy. The live assertions it protects are expensive to reproduce
(they need an upstream that fails the first attempt), which is exactly why the
helper is worth pinning here.
"""

from __future__ import annotations

import pytest
from otel_client import TTFT_TAG, JaegerTrace, one_served_genai_span, served_genai_spans

GENAI_SPAN = "chat claude-haiku-4-5"


def _span(name: str, *, failed: bool = False, ttft: float | None = None) -> dict[str, object]:
    tags: list[dict[str, object]] = []
    if failed:
        tags.append({"key": "otel.status_code", "value": "ERROR"})
        tags.append({"key": "error.type", "value": "AuthenticationError"})
    if ttft is not None:
        tags.append({"key": TTFT_TAG, "value": ttft})
    return {"spanID": f"{name}-{len(tags)}-{failed}-{ttft}", "operationName": name, "tags": tags}


def _trace(*spans: dict[str, object]) -> JaegerTrace:
    return JaegerTrace.model_validate({"traceID": "t1", "spans": list(spans)})


def test_served_span_is_the_only_one_when_nothing_was_retried() -> None:
    trace = _trace(_span("POST /chat/completions"), _span(GENAI_SPAN, ttft=0.3))

    assert [span.operation_name for span in served_genai_spans(trace, GENAI_SPAN)] == [GENAI_SPAN]


def test_retried_attempt_span_is_excluded() -> None:
    """The real shape from a stage trace: the first attempt 401s and records no
    TTFT, the retry serves the stream. The served attempt is the one the TTFT
    assertions must run against."""
    trace = _trace(
        _span(GENAI_SPAN, failed=True),
        _span(GENAI_SPAN, ttft=0.52),
    )

    served = one_served_genai_span(trace, GENAI_SPAN)

    assert [tag.value for tag in served.tags if tag.key == TTFT_TAG] == [0.52]


def test_several_failed_attempts_still_leave_one_served_span() -> None:
    trace = _trace(
        _span(GENAI_SPAN, failed=True),
        _span(GENAI_SPAN, failed=True),
        _span(GENAI_SPAN, failed=True),
        _span(GENAI_SPAN, ttft=0.1),
    )

    assert len(served_genai_spans(trace, GENAI_SPAN)) == 1


def test_two_served_spans_still_fail() -> None:
    """The regression the count assertion exists for: one streamed call must
    not be logged as two served gen-AI spans."""
    trace = _trace(_span(GENAI_SPAN, ttft=0.2), _span(GENAI_SPAN, ttft=0.4))

    with pytest.raises(AssertionError, match="exactly ONE served gen-AI span, got 2"):
        one_served_genai_span(trace, GENAI_SPAN)


def test_all_attempts_failed_is_a_failure_not_a_pass() -> None:
    trace = _trace(_span(GENAI_SPAN, failed=True), _span(GENAI_SPAN, failed=True))

    with pytest.raises(AssertionError, match="exactly ONE served gen-AI span, got 0"):
        one_served_genai_span(trace, GENAI_SPAN)


def test_other_operations_are_not_counted() -> None:
    trace = _trace(_span("chat gpt-5.5", ttft=0.3), _span(GENAI_SPAN, ttft=0.3))

    assert len(served_genai_spans(trace, GENAI_SPAN)) == 1