"""
Tracing utilities for the HIEROMEM API.

Provides optional OpenTelemetry integration with graceful degradation to
no-op tracing when the dependency is not available.
"""

from __future__ import annotations

import logging
import atexit
from contextlib import AbstractContextManager, nullcontext
from functools import wraps
from typing import Any, Awaitable, Callable, Optional, ParamSpec, Protocol, TypeVar, cast

from fastapi import Request, Response


class SpanLike(Protocol):
    """Minimal span protocol used by tracing helpers."""

    def set_attribute(self, name: str, value: Any) -> None:
        ...

    def record_exception(self, exception: BaseException) -> None:
        ...

    def end(self) -> None:
        ...


class TracerLike(Protocol):
    """Protocol describing the subset of tracer behaviour we rely on."""

    def start_as_current_span(
            self,
            name: str) -> AbstractContextManager[SpanLike]:
        ...

    def start_span(self, name: str) -> SpanLike:
        ...


class _NoOpSpan:
    def set_attribute(self, *_: Any, **__: Any) -> None:
        return None

    def record_exception(self, *_: Any, **__: Any) -> None:
        return None

    def end(self) -> None:
        return None


class _NoOpTracer:
    def start_as_current_span(
            self,
            *_: Any,
            **__: Any) -> AbstractContextManager[SpanLike]:
        return nullcontext(_NoOpSpan())

    def start_span(self, *_: Any, **__: Any) -> SpanLike:
        return _NoOpSpan()


class _NoOpTraceAPI:
    def get_tracer(self, *_: Any, **__: Any) -> TracerLike:
        return _NoOpTracer()

    def set_tracer_provider(self, *_: Any, **__: Any) -> None:
        return None


trace: Any = _NoOpTraceAPI()
Resource: Any = None
TracerProvider: Any = None
BatchSpanProcessor: Any = None
ConsoleSpanExporter: Any = None
SimpleSpanProcessor: Any = None
OTLPSpanExporter: Any | None
_HAS_OTEL = False
_ACTIVE_PROVIDER: Any = None
_REGISTERED_SHUTDOWN = False

try:  # pragma: no cover - optional dependency handling
    from opentelemetry import trace as _otel_trace
    from opentelemetry.sdk.resources import Resource as _Resource
    from opentelemetry.sdk.trace import TracerProvider as _TracerProvider
    from opentelemetry.sdk.trace.export import (
        BatchSpanProcessor as _BatchSpanProcessor,
        ConsoleSpanExporter as _ConsoleSpanExporter,
        SimpleSpanProcessor as _SimpleSpanProcessor,
    )
    from opentelemetry.trace import Span as _Span, Tracer as _Tracer

    trace = _otel_trace
    Resource = _Resource
    TracerProvider = _TracerProvider
    BatchSpanProcessor = _BatchSpanProcessor
    ConsoleSpanExporter = _ConsoleSpanExporter
    SimpleSpanProcessor = _SimpleSpanProcessor
    TracerRuntime = _Tracer
    SpanRuntime = _Span
    _HAS_OTEL = True

    try:
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
            OTLPSpanExporter as _OTLPSpanExporter,
        )
        OTLPSpanExporter = _OTLPSpanExporter
    except ImportError:  # pragma: no cover - optional exporter
        OTLPSpanExporter = None
except ImportError:  # pragma: no cover - fallback to no-op tracing
    TracerRuntime = _NoOpTracer
    SpanRuntime = _NoOpSpan
    OTLPSpanExporter = None
else:
    # When OpenTelemetry is installed, ensure runtime types satisfy protocols.
    pass

P = ParamSpec("P")
R = TypeVar("R")

CallNext = Callable[[Request], Awaitable[Response]]


def setup_tracing(
    service_name: str = "hieromem_api",
    exporter: str = "console",
    endpoint: Optional[str] = None,
) -> TracerLike:
    """Initialise global tracer provider if OpenTelemetry is available."""

    logger = logging.getLogger("tracing")

    if not _HAS_OTEL or Resource is None or TracerProvider is None:
        logger.warning(
            "[tracing] OpenTelemetry not installed; returning no-op tracer for %s",
            service_name,
        )
        return cast(TracerLike, trace.get_tracer(service_name))

    resource = Resource.create({"service.name": service_name})
    provider = TracerProvider(resource=resource)

    use_batch_processor = exporter == "otlp" and OTLPSpanExporter
    if use_batch_processor:
        exp = OTLPSpanExporter(
            endpoint=endpoint or "http://localhost:4318/v1/traces")
        logger.info(
            "[tracing] Using OTLP exporter at %s",
            endpoint or "localhost:4318")
    else:
        exp = ConsoleSpanExporter()
        logger.info("[tracing] Using ConsoleSpanExporter (stdout)")

    processor_cls = BatchSpanProcessor if use_batch_processor else SimpleSpanProcessor
    processor = processor_cls(exp)
    provider.add_span_processor(processor)
    trace.set_tracer_provider(provider)
    global _ACTIVE_PROVIDER
    _ACTIVE_PROVIDER = provider
    global _REGISTERED_SHUTDOWN
    if not _REGISTERED_SHUTDOWN:
        atexit.register(shutdown_tracing)
        _REGISTERED_SHUTDOWN = True

    tracer = trace.get_tracer(service_name)
    logger.info(
        "[tracing] OpenTelemetry tracer initialized for %s",
        service_name)
    return cast(TracerLike, tracer)


def shutdown_tracing() -> None:
    """Flush and shutdown the active tracer provider if configured."""

    global _ACTIVE_PROVIDER
    provider = _ACTIVE_PROVIDER
    if not provider:
        return

    try:
        provider.shutdown()
    except Exception as exc:  # pragma: no cover - defensive shutdown
        logging.getLogger("tracing").warning(
            "[tracing] Failed to shutdown tracer provider cleanly: %s", exc
        )
    finally:
        _ACTIVE_PROVIDER = None


async def trace_request(
    tracer: TracerLike | None,
    request: Request,
    call_next: CallNext,
) -> Response:
    """Middleware helper to create a distributed trace span per HTTP request."""

    active_tracer = tracer or cast(TracerLike, trace.get_tracer("hieromem_api"))
    span_name = f"{request.method} {request.url.path}"

    with active_tracer.start_as_current_span(span_name) as span:
        if request.client:
            span.set_attribute("net.peer.ip", request.client.host)
        span.set_attribute("http.method", request.method)
        span.set_attribute("http.target", request.url.path)
        span.set_attribute("http.scheme", request.url.scheme)
        span.set_attribute("http.host", request.url.hostname or "")
        span.set_attribute("request_id", getattr(request.state, "request_id", ""))

        try:
            response = await call_next(request)
            span.set_attribute("http.status_code", getattr(response, "status_code", 0))
            return response
        except Exception as exc:  # pragma: no cover - passthrough to FastAPI handlers
            span.record_exception(exc)
            span.set_attribute("error", True)
            raise


class TracedSection:
    """Context manager for manual tracing of code blocks."""

    def __init__(self, tracer: TracerLike, name: str):
        self.tracer = tracer
        self.name = name
        self.span: SpanLike | None = None

    def __enter__(self) -> SpanLike:
        span = self.tracer.start_span(self.name)
        span.set_attribute("component", "hieromem_core")
        self.span = span
        return span

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None:
        span = self.span
        if span is None:
            return
        if exc_type:
            if exc_val is not None:
                span.record_exception(exc_val)
            span.set_attribute("error", True)
        span.end()


def trace_function(
    tracer: TracerLike | None,
    name: Optional[str] = None,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """Decorator variant of :class:`TracedSection` for function-level tracing."""

    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            span_name = name or func.__name__
            active_tracer = tracer or cast(
                TracerLike, trace.get_tracer(
                    func.__module__))
            with active_tracer.start_as_current_span(span_name) as span:
                try:
                    result = func(*args, **kwargs)
                    span.set_attribute("status", "ok")
                    return result
                except Exception as exc:
                    span.record_exception(exc)
                    span.set_attribute("error", True)
                    raise

        return wrapper

    return decorator
