"""OpenTelemetry tracing setup and instrumentation helpers."""

from typing import cast

from fastapi import FastAPI
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.trace import Tracer, TracerProvider as TraceAPIProvider

_TRACER_PROVIDER: TraceAPIProvider | None = None
_HTTPX_INSTRUMENTED = False


def setup_tracing(
    *,
    service_name: str,
    otlp_endpoint: str,
    tracer_provider: TraceAPIProvider | None = None,
) -> TraceAPIProvider:
    global _TRACER_PROVIDER

    if tracer_provider is None and _TRACER_PROVIDER is not None:
        return _TRACER_PROVIDER

    if tracer_provider is None:
        tracer_provider = TracerProvider(
            resource=Resource.create({"service.name": service_name}),
        )
        if otlp_endpoint:
            span_exporter = OTLPSpanExporter(
                endpoint=otlp_endpoint,
                insecure=otlp_endpoint.startswith("http://"),
            )
            tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))

    _TRACER_PROVIDER = tracer_provider
    return tracer_provider


def instrument_fastapi_app(app: FastAPI) -> None:
    tracer_provider = _TRACER_PROVIDER
    if tracer_provider is None:
        return
    FastAPIInstrumentor.instrument_app(app, tracer_provider=tracer_provider)


def instrument_httpx_client() -> None:
    global _HTTPX_INSTRUMENTED

    tracer_provider = _TRACER_PROVIDER
    if tracer_provider is None:
        return

    if _HTTPX_INSTRUMENTED:
        return
    HTTPXClientInstrumentor().instrument(tracer_provider=tracer_provider)
    _HTTPX_INSTRUMENTED = True


def get_tracer(name: str) -> Tracer:
    tracer_provider = _TRACER_PROVIDER
    if tracer_provider is None:
        return trace.get_tracer(name)
    return cast(Tracer, tracer_provider.get_tracer(name))


def reset_tracing_for_tests() -> None:
    global _TRACER_PROVIDER, _HTTPX_INSTRUMENTED

    if _HTTPX_INSTRUMENTED:
        HTTPXClientInstrumentor().uninstrument()
    if isinstance(_TRACER_PROVIDER, TracerProvider):
        _TRACER_PROVIDER.shutdown()
    _TRACER_PROVIDER = None
    _HTTPX_INSTRUMENTED = False