import wrapt
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.flask import FlaskInstrumentor
from opentelemetry.instrumentation.requests import RequestsInstrumentor
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

from src.internals.tracing.custom_psycopg_instrumentor import instrument_psycopg
from src.internals.tracing.custom_rb_instrumentor import rb_instrument


def open_telemetry_init(app, endpoint):
    resource = Resource.create({"service.name": "kemono3"})
    otlp_exporter = OTLPSpanExporter(
        endpoint=endpoint,
        insecure=True,
    )
    trace.set_tracer_provider(TracerProvider(resource=resource))
    trace.get_tracer_provider().add_span_processor(
        BatchSpanProcessor(
            otlp_exporter,
            schedule_delay_millis=1 * 1e3,  # should be 10
            max_queue_size=15000,
            max_export_batch_size=15000,
        )
    )

    instrument_psycopg()
    FlaskInstrumentor().instrument_app(app)
    RequestsInstrumentor().instrument()
    # RedisInstrumentor().instrument()
    rb_instrument()

    import flask

    def trace_render_template(tracer_name):
        @wrapt.decorator
        def _traced_render_template(func, instance, args, kwargs):
            with trace.get_tracer(tracer_name).start_as_current_span(func.__name__, kind=trace.SpanKind.CLIENT) as span:
                current_span = trace.get_current_span()
                if current_span:
                    current_span.set_attribute("template_name", args[0])
                return func(*args, **kwargs)

        return _traced_render_template

    flask.render_template = trace_render_template("render_template")(flask.render_template)

    import src.internals.database.database

    cached_query_arg_names = src.internals.database.database.cached_query.__annotations__
    cached_query_arg_names.pop("return")
    cached_query_arg_names = tuple(cached_query_arg_names)

    def instrumented_def_cached_query(func, instance, args, kwargs):
        current_span = trace.get_current_span()

        if current_span:
            for arg_name, arg_value in zip(cached_query_arg_names, args):
                current_span.set_attribute(arg_name, arg_value)
            for k, v in kwargs.items():
                current_span.set_attribute(k, v)
            with trace.get_tracer("cached_query").start_as_current_span(
                "cached_query " + args[0].replace("  ", "")[:50], kind=trace.SpanKind.INTERNAL
            ) as span:
                return func(*args, **kwargs)
        return func(*args, **kwargs)

    src.internals.database.database.cached_query = wrapt.decorator(instrumented_def_cached_query)(
        src.internals.database.database.cached_query
    )
    # src.internals.database.database.cached_query = generic_trace("cached_query")(src.internals.database.database.cached_query)
    # wrap_function_wrapper("src", f"internals.database.database.cached_query", generic_trace())
