"""Request correlation, OpenTelemetry spans and stage timing.

This module deliberately has a no-op default. The RAG domain remains runnable
without the OpenTelemetry packages or collector, while a configured runtime
gets one trace context shared by FastAPI's async middleware and its sync
thread-pool endpoint.
"""
from __future__ import annotations

import base64
import re
import uuid
from contextlib import contextmanager
from contextvars import ContextVar
from functools import wraps
from time import perf_counter
from typing import Any, Iterator, Mapping

from .metrics import Metrics, NullMetrics, STAGE_DURATION

_CORRELATION_ID: ContextVar[str | None] = ContextVar(
    "duocthu_correlation_id", default=None
)
_STAGE: ContextVar[str | None] = ContextVar("duocthu_stage", default=None)
_SAFE_CORRELATION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")
_metrics: Metrics = NullMetrics()

try:
    from opentelemetry import trace

    _tracer = trace.get_tracer("duocthu.ai-service")
except ImportError:  # pragma: no cover - exercised only in minimal installs
    trace = None
    _tracer = None


def configure_telemetry(settings: Any, metrics: Metrics | None = None) -> bool:
    """Configure one OTLP tracer provider; return whether tracing is active.

    Metrics are configured independently and continue to work when tracing is
    disabled or optional OpenTelemetry packages are absent.
    """
    global _metrics, _tracer
    _metrics = metrics or NullMetrics()
    if not getattr(settings, "otel_enabled", False):
        return False

    try:
        from opentelemetry import trace as otel_trace
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
        from opentelemetry.sdk.resources import Resource
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import BatchSpanProcessor
        from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased
    except ImportError:
        return False

    current = otel_trace.get_tracer_provider()
    # OpenTelemetry only permits setting the global provider once. Reuse a
    # provider installed by a host/test instead of replacing it and warning.
    if isinstance(current, TracerProvider):
        _tracer = current.get_tracer("duocthu.ai-service")
        return True

    provider = TracerProvider(
        resource=Resource.create(
            {
                "service.name": getattr(settings, "otel_service_name", "ai-service"),
                "service.version": "0.1.0",
                "deployment.environment": getattr(settings, "environment", "local"),
            }
        ),
        sampler=ParentBased(
            TraceIdRatioBased(float(getattr(settings, "otel_sample_ratio", 1.0)))
        ),
    )
    exporter = OTLPSpanExporter(endpoint=settings.otel_exporter_otlp_endpoint)
    provider.add_span_processor(BatchSpanProcessor(exporter))
    for extra in _langfuse_exporters(settings, OTLPSpanExporter):
        provider.add_span_processor(BatchSpanProcessor(extra))
    otel_trace.set_tracer_provider(provider)
    _tracer = provider.get_tracer("duocthu.ai-service")
    return True


def _langfuse_exporters(settings: Any, exporter_cls: Any) -> list[Any]:
    """Zero or one extra OTLP exporter pointed at Langfuse.

    Langfuse accepts standard OTLP/HTTP on `/api/public/otel/v1/traces`,
    authenticated with HTTP Basic using the project's public key as the
    username and the secret key as the password. Returning a list keeps the
    caller unchanged when Langfuse is not configured -- which is the default,
    and the case for every local run without a `.env`.
    """
    base_url = str(getattr(settings, "langfuse_base_url", "") or "").rstrip("/")
    public_key = str(getattr(settings, "langfuse_public_key", "") or "")
    secret_key = str(getattr(settings, "langfuse_secret_key", "") or "")
    if not (base_url and public_key and secret_key):
        return []

    token = base64.b64encode(f"{public_key}:{secret_key}".encode()).decode()
    return [
        exporter_cls(
            endpoint=f"{base_url}/api/public/otel/v1/traces",
            headers={"Authorization": f"Basic {token}"},
        )
    ]


def normalize_correlation_id(candidate: str | None) -> str:
    value = (candidate or "").strip()
    return value if _SAFE_CORRELATION.fullmatch(value) else str(uuid.uuid4())


@contextmanager
def correlation_context(candidate: str | None) -> Iterator[str]:
    correlation_id = normalize_correlation_id(candidate)
    token = _CORRELATION_ID.set(correlation_id)
    try:
        yield correlation_id
    finally:
        _CORRELATION_ID.reset(token)


def current_correlation_id() -> str:
    value = _CORRELATION_ID.get()
    return value or str(uuid.uuid4())


def current_trace_id() -> str | None:
    if trace is None:
        return None
    context = trace.get_current_span().get_span_context()
    if not context.is_valid:
        return None
    return f"{context.trace_id:032x}"


def current_stage() -> str | None:
    return _STAGE.get()


@contextmanager
def request_span(
    method: str, route: str, headers: Mapping[str, str]
) -> Iterator[Any]:
    """Start a server span and extract an incoming W3C ``traceparent``."""
    if _tracer is None:
        yield None
        return
    try:
        from opentelemetry import propagate
        from opentelemetry.trace import SpanKind
    except ImportError:  # pragma: no cover - minimal install fallback
        yield None
        return

    parent = propagate.extract(headers)
    with _tracer.start_as_current_span(
        f"{method} {route}",
        context=parent,
        kind=SpanKind.SERVER,
        attributes={
            "http.request.method": method,
            "http.route": route,
            "duocthu.correlation_id": current_correlation_id(),
        },
    ) as span:
        yield span


@contextmanager
def dependency_span(provider: str, operation: str) -> Iterator[Any]:
    """Trace one external dependency call without creating another histogram."""
    if _tracer is None:
        yield None
        return
    with _tracer.start_as_current_span(
        f"provider.{provider}.{operation}",
        attributes={
            "duocthu.provider": provider,
            "duocthu.operation": operation,
        },
    ) as span:
        yield span


@contextmanager
def stage(name: str, **attributes: Any) -> Iterator[Any]:
    """Create a child span and observe a low-cardinality stage histogram."""
    started = perf_counter()
    outcome = "ok"
    span = None
    token = _STAGE.set(name)
    try:
        if _tracer is None:
            yield None
        else:
            safe_attributes = {
                key: value
                for key, value in attributes.items()
                if value is not None and isinstance(value, (str, bool, int, float))
            }
            safe_attributes["duocthu.stage"] = name
            with _tracer.start_as_current_span(
                f"rag.stage.{name}", attributes=safe_attributes
            ) as span:
                yield span
    except BaseException as exc:
        outcome = "cancelled" if exc.__class__.__name__ == "CancelledError" else "error"
        if span is not None:
            span.set_attribute("duocthu.outcome", outcome)
        raise
    finally:
        elapsed = perf_counter() - started
        if span is not None:
            span.set_attribute("duocthu.outcome", outcome)
            span.set_attribute("duocthu.duration_ms", elapsed * 1000)
        _metrics.observe(STAGE_DURATION, elapsed, stage=name, outcome=outcome)
        _STAGE.reset(token)


def annotate_current_span(**attributes: Any) -> None:
    if trace is None:
        return
    span = trace.get_current_span()
    if not span.is_recording():
        return
    for key, value in attributes.items():
        if value is not None and isinstance(value, (str, bool, int, float)):
            span.set_attribute(key, value)


def traced_stage(name: str):
    """Decorator form for pure functions that should remain otherwise untouched."""
    def decorate(function):
        @wraps(function)
        def wrapped(*args, **kwargs):
            with stage(name):
                return function(*args, **kwargs)

        return wrapped

    return decorate