"""
Tracing utilities for the psycopg potgres client library.
"""

# stdlib
import functools

# 3p
from psycopg2.extensions import connection
from psycopg2.extensions import cursor

from ...constants import SPAN_MEASURED_KEY
from ...ext import SpanTypes
from ...ext import db
from ...ext import net
from ...ext import sql
from ...utils.deprecation import deprecated


@deprecated(message="Use patching instead (see the docs).", version="1.0.0")
def connection_factory(tracer, service="postgres"):
    """Return a connection factory class that will can be used to trace
    postgres queries.

    >>> factory = connection_factor(my_tracer, service='my_db_service')
    >>> conn = pyscopg2.connect(..., connection_factory=factory)
    """

    return functools.partial(
        TracedConnection,
        datadog_tracer=tracer,
        datadog_service=service,
    )


class TracedCursor(cursor):
    """Wrapper around cursor creating one span per query"""

    def __init__(self, *args, **kwargs):
        self._datadog_tracer = kwargs.pop("datadog_tracer", None)
        self._datadog_service = kwargs.pop("datadog_service", None)
        self._datadog_tags = kwargs.pop("datadog_tags", None)
        super(TracedCursor, self).__init__(*args, **kwargs)

    def execute(self, query, vars=None):  # noqa: A002
        """just wrap the cursor execution in a span"""
        if not self._datadog_tracer:
            return cursor.execute(self, query, vars)

        with self._datadog_tracer.trace("postgres.query", service=self._datadog_service, span_type=SpanTypes.SQL) as s:
            s.set_tag(SPAN_MEASURED_KEY)
            if not s.sampled:
                return super(TracedCursor, self).execute(query, vars)

            s.resource = query
            s.set_tags(self._datadog_tags)
            try:
                return super(TracedCursor, self).execute(query, vars)
            finally:
                s.set_metric("db.rowcount", self.rowcount)

    def callproc(self, procname, vars=None):  # noqa: A002
        """just wrap the execution in a span"""
        return cursor.callproc(self, procname, vars)


class TracedConnection(connection):
    """Wrapper around psycopg2  for tracing"""

    def __init__(self, *args, **kwargs):

        self._datadog_tracer = kwargs.pop("datadog_tracer", None)
        self._datadog_service = kwargs.pop("datadog_service", None)

        super(TracedConnection, self).__init__(*args, **kwargs)

        # add metadata (from the connection, string, etc)
        dsn = sql.parse_pg_dsn(self.dsn)
        self._datadog_tags = {
            net.TARGET_HOST: dsn.get("host"),
            net.TARGET_PORT: dsn.get("port"),
            db.NAME: dsn.get("dbname"),
            db.USER: dsn.get("user"),
            "db.application": dsn.get("application_name"),
        }

        self._datadog_cursor_class = functools.partial(
            TracedCursor,
            datadog_tracer=self._datadog_tracer,
            datadog_service=self._datadog_service,
            datadog_tags=self._datadog_tags,
        )

    def cursor(self, *args, **kwargs):
        """register our custom cursor factory"""
        kwargs.setdefault("cursor_factory", self._datadog_cursor_class)
        return super(TracedConnection, self).cursor(*args, **kwargs)
