from __future__ import annotations

from typing import TYPE_CHECKING, Any

from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.redis import RedisInstrumentor
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

from src.infrastructure.config.settings import settings

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncEngine


def setup_tracing(app: Any = None, engine: AsyncEngine | None = None) -> None:
    if not settings.OTEL_TRACES_ENABLED:
        return

    resource = Resource.create({"service.name": settings.OTEL_SERVICE_NAME})

    provider = TracerProvider(resource=resource)

    otlp_exporter = OTLPSpanExporter(endpoint=settings.OTEL_EXPORTER_OTLP_ENDPOINT)
    provider.add_span_processor(BatchSpanProcessor(otlp_exporter))

    trace.set_tracer_provider(provider)

    if engine is not None:
        sync_engine = engine.sync_engine if hasattr(engine, "sync_engine") else engine
        SQLAlchemyInstrumentor().instrument(engine=sync_engine)
    RedisInstrumentor().instrument()

    if app is not None:
        FastAPIInstrumentor.instrument_app(app)


def get_tracer(name: str = "pip-platform") -> trace.Tracer:
    return trace.get_tracer(name)