"""OpenTelemetry W3C trace context propagation via Kafka headers."""

from __future__ import annotations

from typing import Optional

from neonlink.record import HEADER_TRACEPARENT, HEADER_TRACESTATE, Record

try:
    from opentelemetry import context as otel_context
    from opentelemetry import trace
    from opentelemetry.trace import StatusCode
    from opentelemetry.trace.propagation import TraceContextTextMapPropagator

    _HAS_OTEL = True
    _propagator = TraceContextTextMapPropagator()
    _consumer_tracer = trace.get_tracer("neonlink.kafka.consumer")
    _producer_tracer = trace.get_tracer("neonlink.kafka.producer")
except ImportError:
    _HAS_OTEL = False
    _propagator = None
    _consumer_tracer = None
    _producer_tracer = None


class _RecordCarrier:
    """Adapter between Record headers and OTel TextMapPropagator."""

    def __init__(self, record: Record) -> None:
        self._record = record

    def get(self, key: str) -> Optional[str]:
        val = self._record.headers.get(key)
        if val is None:
            return None
        if isinstance(val, bytes):
            return val.decode("utf-8", errors="replace")
        return val or None

    def set(self, key: str, value: str) -> None:
        self._record.headers[key] = value.encode("utf-8")

    def keys(self) -> list[str]:
        return list(self._record.headers.keys())


def inject_trace_context(record: Record) -> None:
    """Inject the current span's W3C trace context into record headers.

    Call this before producing a record. If OpenTelemetry is not installed,
    this is a no-op.
    """
    if not _HAS_OTEL:
        return
    _propagator.inject(carrier=_RecordCarrier(record))


def extract_trace_context(record: Record) -> Optional[object]:
    """Extract W3C trace context from record headers.

    Returns an OTel Context object with the remote span context attached.
    If OpenTelemetry is not installed, returns None.
    """
    if not _HAS_OTEL:
        return None
    return _propagator.extract(carrier=_RecordCarrier(record))


class _NoOpSpan:
    """Minimal no-op span for when OTel is not installed."""

    def set_status(self, *args, **kwargs):
        pass

    def record_error(self, *args, **kwargs):
        pass

    def end(self):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass


_NOOP_SPAN = _NoOpSpan()


def start_consumer_span(record: Record):
    """Start a consumer span for a consumed record.

    Returns (span,). Callers must call span.end() when processing is complete.
    If OTel is not installed, returns a no-op span.
    """
    if not _HAS_OTEL:
        return _NOOP_SPAN
    return _consumer_tracer.start_span(
        f"{record.topic} receive",
        kind=trace.SpanKind.CONSUMER,
        attributes={
            "messaging.system": "kafka",
            "messaging.destination": record.topic,
            "messaging.kafka.offset": record.offset,
            "messaging.kafka.partition": record.partition,
        },
    )


def start_producer_span(topic: str):
    """Start a producer span for a produced record.

    Returns (span,). Callers must call span.end() when publishing is complete.
    If OTel is not installed, returns a no-op span.
    """
    if not _HAS_OTEL:
        return _NOOP_SPAN
    return _producer_tracer.start_span(
        f"{topic} send",
        kind=trace.SpanKind.PRODUCER,
        attributes={
            "messaging.system": "kafka",
            "messaging.destination": topic,
        },
    )


def end_span_with_error(span, err: Optional[Exception] = None) -> None:
    """End a span, optionally recording an error."""
    if err is not None and _HAS_OTEL:
        span.record_error(err)
        span.set_status(StatusCode.ERROR, str(err))
    span.end()
