import re
import typing

import wrapt
from opentelemetry import trace
from opentelemetry.semconv.trace import SpanAttributes
from opentelemetry.trace import Span, SpanKind
from psycopg import Cursor
from psycopg.sql import Composed

_leading_comment_remover = re.compile(r"^/\*.*?\*/")
tracer = trace.get_tracer("psycopg3", "0")

instrumented = False


def instrument_psycopg():
    global instrumented
    if instrumented:
        return
    instrumented = True

    def _populate_span(
        span: Span,
        cursor,
        *args: typing.Tuple[typing.Any, typing.Any],
    ):
        if not span.is_recording():
            return
        statement = get_statement(cursor, args)
        span.set_attribute(SpanAttributes.DB_SYSTEM, cursor._conn.pgconn.db)

        span.set_attribute(SpanAttributes.DB_STATEMENT, statement)

        span.set_attribute("db.statement.parameters", str(args[1]) if len(args) > 1 else "no parameters")

    def get_operation_name(cursor, args):
        if not args:
            return ""

        statement = args[0]
        if isinstance(statement, Composed):
            statement = statement.as_string(cursor)

        if isinstance(statement, str):
            # Strip leading comments so we get the operation name.
            return _leading_comment_remover.sub("", statement).split()[0]

        return ""

    def get_statement(cursor, args):
        if not args:
            return ""

        statement = args[0]
        if isinstance(statement, Composed):
            statement = statement.as_string(cursor)
        return statement

    @wrapt.decorator
    def traced_execution(func, instance, args, kwargs):
        name = get_operation_name(instance, args)
        if not name:
            name = "traced_execution"

        with tracer.start_as_current_span(name, kind=SpanKind.CLIENT) as span:
            _populate_span(span, instance, *args)
            return func(*args, **kwargs)

    Cursor.execute = traced_execution(Cursor.execute)
    Cursor.executemany = traced_execution(Cursor.executemany)

    @wrapt.decorator
    def count_results(func, instance, args, kwargs):
        with trace.get_tracer("fetch_all").start_as_current_span("fetchall", kind=trace.SpanKind.CLIENT) as span:
            response = func(*args, **kwargs)
            span.set_attribute("len_results", len(response))
            return response

    Cursor.fetchall = count_results(Cursor.fetchall)
