from ddtrace import config
from ddtrace._trace.pin import Pin
from ddtrace.constants import SPAN_KIND
from ddtrace.contrib import dbapi_async
from ddtrace.contrib.internal.psycopg.async_cursor import Psycopg3FetchTracedAsyncCursor
from ddtrace.contrib.internal.psycopg.async_cursor import Psycopg3TracedAsyncCursor
from ddtrace.contrib.internal.psycopg.connection import patch_conn
from ddtrace.contrib.internal.trace_utils import ext_service
from ddtrace.ext import SpanKind
from ddtrace.ext import SpanTypes
from ddtrace.ext import db
from ddtrace.internal import core
from ddtrace.internal.constants import COMPONENT


class Psycopg3TracedAsyncConnection(dbapi_async.TracedAsyncConnection):
    def __init__(self, conn, pin=None, cursor_cls=None):
        if not cursor_cls:
            # Do not trace `fetch*` methods by default
            cursor_cls = (
                Psycopg3FetchTracedAsyncCursor if config.psycopg.trace_fetch_methods else Psycopg3TracedAsyncCursor
            )

        super(Psycopg3TracedAsyncConnection, self).__init__(conn, pin, config.psycopg, cursor_cls=cursor_cls)

    async def execute(self, *args, **kwargs):
        """Execute a query and return a cursor to read its results."""
        span_name = "{}.{}".format(self._self_datadog_name, "execute")

        async def patched_execute(*args, **kwargs):
            try:
                cur = self.cursor()
                if kwargs.get("binary", None):
                    cur.format = 1  # set to 1 for binary or 0 if not
                return await cur.execute(*args, **kwargs)
            except Exception as ex:
                raise ex.with_traceback(None)

        return await self._trace_method(patched_execute, span_name, {}, *args, **kwargs)


def patched_connect_async_factory(psycopg_module):
    async def patched_connect_async(connect_func, _, args, kwargs):
        traced_conn_cls = Psycopg3TracedAsyncConnection

        pin = Pin.get_from(psycopg_module)

        if not pin or not pin.enabled() or not pin._config.trace_connect:
            conn = await connect_func(*args, **kwargs)
        else:
            with core.context_with_data(
                "psycopg.patched_connect",
                span_name="{}.{}".format(connect_func.__module__, connect_func.__name__),
                service=ext_service(pin, pin._config),
                span_type=SpanTypes.SQL,
                pin=pin,
                tags={
                    SPAN_KIND: SpanKind.CLIENT,
                    COMPONENT: pin._config.integration_name,
                    db.SYSTEM: pin._config.dbms_name,
                },
                measured=True,
            ):
                conn = await connect_func(*args, **kwargs)

        return patch_conn(conn, pin=pin, traced_conn_cls=traced_conn_cls)

    return patched_connect_async
