"""LLMHosts distributed tracing -- optional OpenTelemetry integration.

Provides request-level distributed tracing for the LLMHosts proxy using
OpenTelemetry.  All OTel dependencies are optional -- when not installed,
a lightweight no-op tracer is used so the rest of the codebase can call
``get_tracer()`` unconditionally.

This module is part of T-12 (World-Class Logging & Monitoring) compliance.

Usage::

    from llmhosts.observability.tracing import TracingConfig, init_tracing, get_tracer

    config = TracingConfig(enabled=True, exporter="console")
    init_tracing(config)
    tracer = get_tracer()
    with tracer.start_as_current_span("my-operation"):
        ...
"""

from __future__ import annotations

import logging
import time
from typing import TYPE_CHECKING, Any, Literal

from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from starlette.types import ASGIApp, Receive, Scope, Send

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Optional OpenTelemetry imports
# ---------------------------------------------------------------------------

_HAS_OTEL = False

try:
    from opentelemetry import trace as otel_trace
    from opentelemetry.sdk.resources import Resource
    from opentelemetry.sdk.trace import TracerProvider
    from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter, SimpleSpanProcessor

    _HAS_OTEL = True
except ImportError:
    otel_trace = None  # type: ignore[assignment]
    Resource = None  # type: ignore[assignment,misc]
    TracerProvider = None  # type: ignore[assignment,misc]
    BatchSpanProcessor = None  # type: ignore[assignment,misc]
    ConsoleSpanExporter = None  # type: ignore[assignment,misc]
    SimpleSpanProcessor = None  # type: ignore[assignment,misc]


def _has_otel() -> bool:
    """Return True if OpenTelemetry SDK is importable."""
    return _HAS_OTEL


# ---------------------------------------------------------------------------
# No-op tracer implementation (used when OTel is not installed or disabled)
# ---------------------------------------------------------------------------


class _NoOpSpan:
    """Minimal span stand-in that accepts attribute writes and context-manager usage."""

    def set_attribute(self, key: str, value: Any) -> None:
        """No-op: discard attribute."""

    def set_status(self, status: Any, description: str | None = None) -> None:
        """No-op: discard status."""

    def record_exception(self, exception: BaseException, attributes: dict[str, Any] | None = None) -> None:
        """No-op: discard exception."""

    def end(self) -> None:
        """No-op: nothing to finalize."""

    def __enter__(self) -> _NoOpSpan:
        return self

    def __exit__(self, *args: Any) -> None:
        pass


class _NoOpTracer:
    """Tracer that produces :class:`_NoOpSpan` instances.

    This allows callers to write ``get_tracer().start_as_current_span(...)``
    without guarding every call with an ``if tracing_available`` check.
    """

    def start_as_current_span(self, name: str, **kwargs: Any) -> _NoOpSpan:
        """Return a no-op span context manager."""
        return _NoOpSpan()

    def start_span(self, name: str, **kwargs: Any) -> _NoOpSpan:
        """Return a no-op span."""
        return _NoOpSpan()


# ---------------------------------------------------------------------------
# Module-level state
# ---------------------------------------------------------------------------

_tracer: Any = _NoOpTracer()
_initialized: bool = False

# ---------------------------------------------------------------------------
# Config model
# ---------------------------------------------------------------------------


class TracingConfig(BaseModel):
    """Configuration for OpenTelemetry distributed tracing.

    All fields have safe defaults so the proxy starts without any
    tracing configuration.
    """

    enabled: bool = Field(default=False, description="Enable distributed tracing")
    exporter: Literal["otlp", "jaeger", "console"] = Field(
        default="otlp", description="Span exporter backend (otlp, jaeger, or console)"
    )
    endpoint: str = Field(default="http://localhost:4317", description="Exporter endpoint (OTLP gRPC or Jaeger)")
    service_name: str = Field(default="llmhosts-proxy", description="Service name reported in traces")
    sample_rate: float = Field(default=1.0, ge=0.0, le=1.0, description="Sampling rate (0.0 to 1.0)")


# ---------------------------------------------------------------------------
# Initialization
# ---------------------------------------------------------------------------


def init_tracing(config: TracingConfig) -> None:
    """Set up the OpenTelemetry TracerProvider, exporter, and resource.

    When tracing is disabled or the OTel SDK is not installed, this function
    is a no-op (the module-level tracer remains a :class:`_NoOpTracer`).

    Parameters
    ----------
    config:
        Tracing configuration from the ``[observability.tracing]`` TOML section.
    """
    global _tracer, _initialized

    if not config.enabled:
        logger.debug("Tracing disabled by configuration")
        _tracer = _NoOpTracer()
        _initialized = True
        return

    if not _has_otel():
        logger.warning(
            "OpenTelemetry SDK not installed -- tracing disabled.  Install with: pip install llmhosts[observability]"
        )
        _tracer = _NoOpTracer()
        _initialized = True
        return

    # Build the OTel resource (service identity)
    resource = Resource.create(
        {
            "service.name": config.service_name,
            "service.version": _get_version(),
        }
    )

    # Create the TracerProvider with sampling
    from opentelemetry.sdk.trace import sampling as otel_sampling

    if config.sample_rate >= 1.0:
        sampler = otel_sampling.ALWAYS_ON
    elif config.sample_rate <= 0.0:
        sampler = otel_sampling.ALWAYS_OFF
    else:
        sampler = otel_sampling.TraceIdRatioBased(config.sample_rate)

    provider = TracerProvider(resource=resource, sampler=sampler)

    # Attach the exporter
    exporter = _create_exporter(config)
    if config.exporter == "console":
        # Console exporter uses SimpleSpanProcessor for immediate output
        provider.add_span_processor(SimpleSpanProcessor(exporter))
    else:
        # Network exporters use BatchSpanProcessor for efficiency
        provider.add_span_processor(BatchSpanProcessor(exporter))

    # Register as global TracerProvider
    otel_trace.set_tracer_provider(provider)
    _tracer = provider.get_tracer("llmhosts", _get_version())
    _initialized = True

    logger.info(
        "OpenTelemetry tracing initialized (exporter=%s, endpoint=%s, sample_rate=%.2f)",
        config.exporter,
        config.endpoint,
        config.sample_rate,
    )


def _create_exporter(config: TracingConfig) -> Any:
    """Instantiate the configured span exporter.

    Falls back to :class:`ConsoleSpanExporter` if the requested exporter
    package is not installed.
    """
    if config.exporter == "console":
        return ConsoleSpanExporter()

    if config.exporter == "otlp":
        try:
            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter

            return OTLPSpanExporter(endpoint=config.endpoint)
        except ImportError:
            logger.warning(
                "opentelemetry-exporter-otlp not installed -- falling back to console exporter.  "
                "Install with: pip install opentelemetry-exporter-otlp"
            )
            return ConsoleSpanExporter()

    if config.exporter == "jaeger":
        try:
            from opentelemetry.exporter.jaeger.thrift import JaegerExporter

            return JaegerExporter(agent_host_name=config.endpoint.split("://")[-1].split(":")[0])
        except ImportError:
            logger.warning(
                "opentelemetry-exporter-jaeger not installed -- falling back to console exporter.  "
                "Install with: pip install opentelemetry-exporter-jaeger"
            )
            return ConsoleSpanExporter()

    # Unreachable due to Literal type, but defensive
    logger.warning("Unknown exporter %r -- falling back to console", config.exporter)
    return ConsoleSpanExporter()


def get_tracer() -> Any:
    """Return the current tracer instance.

    If :func:`init_tracing` has not been called or tracing is disabled,
    returns a :class:`_NoOpTracer` so callers never need to guard usage.

    Returns
    -------
    opentelemetry.trace.Tracer | _NoOpTracer
        A tracer that supports ``start_as_current_span`` and ``start_span``.
    """
    return _tracer


def reset_tracing() -> None:
    """Reset module state -- intended for tests only.

    Restores the tracer to a :class:`_NoOpTracer` and clears the
    initialized flag so :func:`init_tracing` can be called again.
    """
    global _tracer, _initialized
    _tracer = _NoOpTracer()
    _initialized = False


# ---------------------------------------------------------------------------
# ASGI Middleware
# ---------------------------------------------------------------------------


class TracingMiddleware:
    """Starlette-compatible ASGI middleware that creates a root span per request.

    When tracing is disabled, the middleware is a zero-cost pass-through --
    it delegates directly to the inner application with no span creation.

    Attributes set on each span:
        - ``http.method`` -- request method (GET, POST, ...)
        - ``http.url`` -- full request URL
        - ``http.route`` -- URL path
        - ``http.status_code`` -- response status code
        - ``http.request_id`` -- X-Request-ID header if present
    """

    def __init__(self, app: ASGIApp, enabled: bool = True) -> None:
        self.app = app
        self.enabled = enabled

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """ASGI entry point."""
        if scope["type"] != "http" or not self.enabled:
            await self.app(scope, receive, send)
            return

        tracer = get_tracer()

        # If we have a no-op tracer, skip span creation entirely for zero overhead
        if isinstance(tracer, _NoOpTracer):
            await self.app(scope, receive, send)
            return

        # Extract request metadata from ASGI scope
        method = scope.get("method", "UNKNOWN")
        path = scope.get("path", "/")
        request_id = ""
        for raw_name, raw_value in scope.get("headers", []):
            if raw_name == b"x-request-id":
                request_id = raw_value.decode("utf-8", errors="replace")
                break

        span_name = f"{method} {path}"
        status_code = 500  # default until we capture the real status

        # Capture response status code from send messages
        async def _send_wrapper(message: Any) -> None:
            nonlocal status_code
            if message.get("type") == "http.response.start":
                status_code = message.get("status", 500)
            await send(message)

        with tracer.start_as_current_span(span_name) as span:
            span.set_attribute("http.method", method)
            span.set_attribute("http.url", _build_url_from_scope(scope))
            span.set_attribute("http.route", path)
            if request_id:
                span.set_attribute("http.request_id", request_id)

            start_time = time.perf_counter()
            try:
                await self.app(scope, receive, _send_wrapper)
            except Exception as exc:
                span.set_attribute("http.status_code", 500)
                span.record_exception(exc)
                if _has_otel():
                    from opentelemetry.trace import StatusCode

                    span.set_status(StatusCode.ERROR, str(exc))
                raise
            else:
                span.set_attribute("http.status_code", status_code)
                if _has_otel() and status_code >= 400:
                    from opentelemetry.trace import StatusCode

                    span.set_status(StatusCode.ERROR, f"HTTP {status_code}")
            finally:
                elapsed_ms = (time.perf_counter() - start_time) * 1000
                span.set_attribute("http.duration_ms", round(elapsed_ms, 2))


def _build_url_from_scope(scope: Scope) -> str:
    """Reconstruct the request URL from an ASGI scope dict."""
    scheme = scope.get("scheme", "http")
    server = scope.get("server")
    path = scope.get("path", "/")
    query_string = scope.get("query_string", b"")

    if server:
        host, port = server
        # Omit default ports
        if (scheme == "http" and port == 80) or (scheme == "https" and port == 443):
            base = f"{scheme}://{host}{path}"
        else:
            base = f"{scheme}://{host}:{port}{path}"
    else:
        base = path

    if query_string:
        return f"{base}?{query_string.decode('utf-8', errors='replace')}"
    return base


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


def _get_version() -> str:
    """Return the llmhost package version, or 'unknown' if not importable."""
    try:
        from llmhosts import __version__

        return __version__
    except ImportError:
        return "unknown"
