"""Flowlines SpanExporter that filters LLM spans and forwards via OTLP/HTTP."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from opentelemetry.exporter.otlp.proto.http import Compression
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
    OTLPSpanExporter,
)
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult

from flowlines._filter import is_llm_span

if TYPE_CHECKING:
    from collections.abc import Sequence

    from opentelemetry.sdk.trace import ReadableSpan

logger = logging.getLogger("flowlines")


class FlowlinesExporter(SpanExporter):
    """A SpanExporter that filters LLM spans and exports via OTLP/HTTP."""

    def __init__(self, api_key: str, endpoint: str) -> None:
        self._inner = OTLPSpanExporter(
            endpoint=f"{endpoint.rstrip('/')}/v1/traces",
            compression=Compression.Gzip,
            headers={"x-flowlines-api-key": api_key},
        )
        self._warned_missing_context: bool = False

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        llm_spans = [s for s in spans if is_llm_span(s)]
        if not llm_spans:
            logger.debug(
                "Export: %d span(s) received, none are LLM-related — skipping",
                len(spans),
            )
            return SpanExportResult.SUCCESS
        if not self._warned_missing_context:
            for span in llm_spans:
                attrs = span.attributes or {}
                has_user = "flowlines.user_id" in attrs
                has_session = "flowlines.session_id" in attrs
                if not has_user or not has_session:
                    logger.warning(
                        "LLM spans are missing flowlines.user_id and/or "
                        "flowlines.session_id attributes. These spans won't be "
                        "processed by the Flowlines backend. Wrap your LLM calls "
                        "with: with flowlines.context(user_id=..., session_id=...)"
                    )
                    self._warned_missing_context = True
                    break
        logger.debug(
            "Export: %d/%d span(s) are LLM-related — sending to backend",
            len(llm_spans),
            len(spans),
        )
        result = self._inner.export(llm_spans)
        if result == SpanExportResult.SUCCESS:
            logger.debug("Export: succeeded")
        else:
            logger.debug("Export: failed with result %s", result)
        return result

    def shutdown(self) -> None:
        self._inner.shutdown()  # type: ignore[no-untyped-call]

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return self._inner.force_flush(timeout_millis)
