import atexit
import contextlib
import logging
import os
import re
import typing

from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.instrumentation.pymongo import PymongoInstrumentor
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
from opentelemetry.sdk.trace.export import (
    BatchSpanProcessor,
    ConsoleSpanExporter,
    SpanExporter,
    SpanExportResult,
)

from delphai_utils.logging import register_custom_field

SERVICE_NAME_RE = re.compile(r"[^a-z0-9-]+")


class UniqueMongoCommandsSpanExporter(SpanExporter):
    IGNORE_STATEMENTS = ["saslContinue", "getMore"]

    def __init__(self):
        self.logger = logging.getLogger("UniqueMongoCommandsSpanExporter")
        self.unique_statements = {}

        atexit.register(self.log_unique_statements)

    def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult:
        server_span_names = {}
        mongo_statements_per_trace = {}

        for span in spans:
            trace_id = span.context.trace_id
            if not trace_id:
                continue

            if span.kind == trace.SpanKind.SERVER:
                server_span_names[trace_id] = span.name
                continue

            if span.kind != trace.SpanKind.CLIENT:
                continue

            scope_name = span.instrumentation_scope.name
            if scope_name != "opentelemetry.instrumentation.pymongo":
                continue

            if span.attributes.get("db.statement") in self.IGNORE_STATEMENTS:
                continue

            mongo_statements_per_trace.setdefault(trace_id, set()).add(
                frozenset(
                    {
                        "database": span.attributes.get("db.name"),
                        "collection": span.attributes.get("db.mongodb.collection"),
                        "statement": span.attributes.get("db.statement"),
                    }.items()
                )
            )

        updated = False
        for trace_id, mongo_requests in mongo_statements_per_trace.items():
            server_request_name = server_span_names.get(trace_id)
            if not server_request_name:
                continue

            unique_statements = self.unique_statements.setdefault(
                server_request_name, set()
            )

            if mongo_requests <= unique_statements:
                continue

            updated = True
            unique_statements |= mongo_requests

        if updated:
            self.log_unique_statements()

        return SpanExportResult.SUCCESS

    def log_unique_statements(self):
        for name, queries in sorted(self.unique_statements.items()):
            for query in queries:
                query = dict(query)
                query = f"{query.get('database')}.{query.get('collection')}: {query.get('statement')}"
                self.logger.info("%s => %s", name, query)


def instrument_mongo():
    def _pymongo_request_hook(span, event):
        if not span.is_recording():
            return

        collection = span.attributes.get("db.mongodb.collection") or "*"
        span_name = f"mongo:{event.database_name}.{collection}.{event.command_name}"

        span.update_name(span_name)

    # Applied once per process
    # Affects only instances created after patch
    PymongoInstrumentor().instrument(request_hook=_pymongo_request_hook)


def configure_tracing(service_name):
    use_exporter_console = os.getenv("OTEL_EXPORTER_CONSOLE")
    use_exporter_otlp = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT") or os.getenv(
        "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"
    )

    if not use_exporter_console and not use_exporter_otlp:
        return

    # Allow to overwrite service name
    service_name = SERVICE_NAME_RE.sub("-", service_name)
    service_name = os.getenv("OTEL_SERVICE_NAME", service_name).strip().lower()
    resource = Resource.create(attributes={SERVICE_NAME: service_name})

    tracer_provider = TracerProvider(resource=resource)

    if use_exporter_console:
        tracer_provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))

    if use_exporter_otlp:
        tracer_provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter()))

    with contextlib.suppress(ImportError):
        import httpx  # noqa: F401

        _instrument_httpx()

    with contextlib.suppress(ImportError):
        import pymongo  # noqa: F401

        tracer_provider.add_span_processor(
            BatchSpanProcessor(UniqueMongoCommandsSpanExporter())
        )

    trace.set_tracer_provider(tracer_provider)

    def get_trace_id():
        span_context = trace.get_current_span().get_span_context()
        if span_context.trace_id:
            trace_id = span_context.trace_id
            span_id = span_context.span_id
            return f"{trace_id:016x}:{span_id:016x}"

    register_custom_field("trace_id", get_trace_id)


def _instrument_httpx():
    def request_hook(span, request):
        span.update_name(f"httpx:{span.name}")

    async def async_request_hook(span, request):
        span.update_name(f"httpx:{span.name}")

    # Applied once per process
    # Patches class methods, affects new and existing instances
    HTTPXClientInstrumentor().instrument(
        request_hook=request_hook,
        async_request_hook=async_request_hook,
    )
