# Copyright (C) 2005-2023 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: https://www.opensource.org/licenses/mit-license.php
# mypy: ignore-errors
# NOTE: This file was copied from https://raw.githubusercontent.com/sqlalchemy/sqlalchemy/rel_2_0_13/lib/sqlalchemy/dialects/postgresql/psycopg.py.
# Chalk modified it to be compatible with sqlalchemy 1.4

r"""
.. dialect:: postgresql+psycopg
    :name: psycopg (a.k.a. psycopg 3)
    :dbapi: psycopg
    :connectstring: postgresql+psycopg://user:password@host:port/dbname[?key=value&key=value...]
    :url: https://pypi.org/project/psycopg/

``psycopg`` is the package and module name for version 3 of the ``psycopg``
database driver, formerly known as ``psycopg2``.  This driver is different
enough from its ``psycopg2`` predecessor that SQLAlchemy supports it
via a totally separate dialect; support for ``psycopg2`` is expected to remain
for as long as that package continues to function for modern Python versions,
and also remains the default dialect for the ``postgresql://`` dialect
series.

The SQLAlchemy ``psycopg`` dialect provides both a sync and an async
implementation under the same dialect name. The proper version is
selected depending on how the engine is created:

* calling :func:`_sa.create_engine` with ``postgresql+psycopg://...`` will
  automatically select the sync version, e.g.::

    from sqlalchemy import create_engine
    sync_engine = create_engine("postgresql+psycopg://scott:tiger@localhost/test")

* calling :func:`_asyncio.create_async_engine` with
  ``postgresql+psycopg://...`` will automatically select the async version,
  e.g.::

    from sqlalchemy.ext.asyncio import create_async_engine
    asyncio_engine = create_async_engine("postgresql+psycopg://scott:tiger@localhost/test")

The asyncio version of the dialect may also be specified explicitly using the
``psycopg_async`` suffix, as::

    from sqlalchemy.ext.asyncio import create_async_engine
    asyncio_engine = create_async_engine("postgresql+psycopg_async://scott:tiger@localhost/test")

The ``psycopg`` dialect has the same API features as that of ``psycopg2``,
with the exception of the "fast executemany" helpers.   The "fast executemany"
helpers are expected to be generalized and ported to ``psycopg`` before the final
release of SQLAlchemy 2.0, however.


.. seealso::

    :ref:`postgresql_psycopg2` - The SQLAlchemy ``psycopg``
    dialect shares most of its behavior with the ``psycopg2`` dialect.
    Further documentation is available there.

"""  # noqa
from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING, cast

from sqlalchemy import pool, util
from sqlalchemy.dialects.postgresql import ranges
from sqlalchemy.dialects.postgresql.base import INTERVAL, PGCompiler, PGIdentifierPreparer

# from sqlalchemy.dialects.postgresql.base import REGCONFIG
from sqlalchemy.dialects.postgresql.json import JSON, JSONB, JSONPathType
from sqlalchemy.engine import AdaptedConnection
from sqlalchemy.sql import sqltypes
from sqlalchemy.util.concurrency import await_fallback, await_only

from chalk.sql._internal.integrations.psycopg3.psycopg_common import (
    _PGDialect_common_psycopg,
    _PGExecutionContext_common_psycopg,
)

if TYPE_CHECKING:
    from typing import Iterable

logger = logging.getLogger("sqlalchemy.dialects.postgresql")


class _PGString(sqltypes.String):
    render_bind_cast = True


# class _PGREGCONFIG(REGCONFIG):
#     render_bind_cast = True


class _PGJSON(JSON):
    render_bind_cast = True

    def bind_processor(self, dialect):
        return self._make_bind_processor(None, dialect._psycopg_Json)

    def result_processor(self, dialect, coltype):
        return None


class _PGJSONB(JSONB):
    render_bind_cast = True

    def bind_processor(self, dialect):
        return self._make_bind_processor(None, dialect._psycopg_Jsonb)

    def result_processor(self, dialect, coltype):
        return None


class _PGJSONIntIndexType(sqltypes.JSON.JSONIntIndexType):
    __visit_name__ = "json_int_index"

    render_bind_cast = True


class _PGJSONStrIndexType(sqltypes.JSON.JSONStrIndexType):
    __visit_name__ = "json_str_index"

    render_bind_cast = True


class _PGJSONPathType(JSONPathType):
    pass


class _PGInterval(INTERVAL):
    render_bind_cast = True


class _PGTimeStamp(sqltypes.DateTime):
    render_bind_cast = True


class _PGDate(sqltypes.Date):
    render_bind_cast = True


class _PGTime(sqltypes.Time):
    render_bind_cast = True


class _PGInteger(sqltypes.Integer):
    render_bind_cast = True


class _PGSmallInteger(sqltypes.SmallInteger):
    render_bind_cast = True


class _PGNullType(sqltypes.NullType):
    render_bind_cast = True


class _PGBigInteger(sqltypes.BigInteger):
    render_bind_cast = True


class _PGBoolean(sqltypes.Boolean):
    render_bind_cast = True


# class _PsycopgRange(ranges.AbstractRangeImpl):
#     def bind_processor(self, dialect):
#         psycopg_Range = cast(PGDialect_psycopg, dialect)._psycopg_Range

#         def to_range(value):
#             if isinstance(value, ranges.Range):
#                 value = psycopg_Range(
#                     value.lower, value.upper, value.bounds, value.empty
#                 )
#             return value

#         return to_range

#     def result_processor(self, dialect, coltype):
#         def to_range(value):
#             if value is not None:
#                 value = ranges.Range(
#                     value._lower,
#                     value._upper,
#                     bounds=value._bounds if value._bounds else "[)",
#                     empty=not value._bounds,
#                 )
#             return value

#         return to_range


# class _PsycopgMultiRange(ranges.AbstractMultiRangeImpl):
#     def bind_processor(self, dialect):
#         psycopg_Range = cast(PGDialect_psycopg, dialect)._psycopg_Range
#         psycopg_Multirange = cast(
#             PGDialect_psycopg, dialect
#         )._psycopg_Multirange

#         NoneType = type(None)

#         def to_range(value):
#             if isinstance(value, (str, NoneType, psycopg_Multirange)):
#                 return value

#             return psycopg_Multirange(
#                 [
#                     psycopg_Range(
#                         element.lower,
#                         element.upper,
#                         element.bounds,
#                         element.empty,
#                     )
#                     for element in cast("Iterable[ranges.Range]", value)
#                 ]
#             )

#         return to_range

#     def result_processor(self, dialect, coltype):
#         def to_range(value):
#             if value is not None:
#                 value = [
#                     ranges.Range(
#                         elem._lower,
#                         elem._upper,
#                         bounds=elem._bounds if elem._bounds else "[)",
#                         empty=not elem._bounds,
#                     )
#                     for elem in value
#                 ]

#             return value

#         return to_range


class PGExecutionContext_psycopg(_PGExecutionContext_common_psycopg):
    pass


class PGCompiler_psycopg(PGCompiler):
    pass


class PGIdentifierPreparer_psycopg(PGIdentifierPreparer):
    pass


def _log_notices(diagnostic):
    logger.info("%s: %s", diagnostic.severity, diagnostic.message_primary)


class PGDialect_psycopg(_PGDialect_common_psycopg):
    driver = "psycopg"

    supports_statement_cache = True
    supports_server_side_cursors = True
    default_paramstyle = "pyformat"
    supports_sane_multi_rowcount = True

    execution_ctx_cls = PGExecutionContext_psycopg
    statement_compiler = PGCompiler_psycopg
    preparer = PGIdentifierPreparer_psycopg
    psycopg_version = (0, 0)

    _has_native_hstore = True
    _psycopg_adapters_map = None

    colspecs = util.update_copy(
        _PGDialect_common_psycopg.colspecs,
        {
            sqltypes.String: _PGString,
            # REGCONFIG: _PGREGCONFIG,
            JSON: _PGJSON,
            sqltypes.JSON: _PGJSON,
            JSONB: _PGJSONB,
            sqltypes.JSON.JSONPathType: _PGJSONPathType,
            sqltypes.JSON.JSONIntIndexType: _PGJSONIntIndexType,
            sqltypes.JSON.JSONStrIndexType: _PGJSONStrIndexType,
            sqltypes.Interval: _PGInterval,
            INTERVAL: _PGInterval,
            sqltypes.Date: _PGDate,
            sqltypes.DateTime: _PGTimeStamp,
            sqltypes.Time: _PGTime,
            sqltypes.Integer: _PGInteger,
            sqltypes.SmallInteger: _PGSmallInteger,
            sqltypes.BigInteger: _PGBigInteger,
            # ranges.AbstractRange: _PsycopgRange,
            # ranges.AbstractMultiRange: _PsycopgMultiRange,
        },
    )

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        if self.dbapi:
            m = re.match(r"(\d+)\.(\d+)(?:\.(\d+))?", self.dbapi.__version__)
            if m:
                self.psycopg_version = tuple(int(x) for x in m.group(1, 2, 3) if x is not None)

            if self.psycopg_version < (3, 0, 2):
                raise ImportError("psycopg version 3.0.2 or higher is required.")

            from psycopg.adapt import AdaptersMap

            self._psycopg_adapters_map = adapters_map = AdaptersMap(self.dbapi.adapters)

            if self._json_deserializer:
                from psycopg.types.json import set_json_loads

                set_json_loads(self._json_deserializer, adapters_map)

            if self._json_serializer:
                from psycopg.types.json import set_json_dumps

                set_json_dumps(self._json_serializer, adapters_map)

    def create_connect_args(self, url):
        # see https://github.com/psycopg/psycopg/issues/83
        cargs, cparams = super().create_connect_args(url)

        if self._psycopg_adapters_map:
            cparams["context"] = self._psycopg_adapters_map
        if self.client_encoding is not None:
            cparams["client_encoding"] = self.client_encoding
        return cargs, cparams

    def _type_info_fetch(self, connection, name):
        from psycopg.types import TypeInfo

        return TypeInfo.fetch(connection.connection.driver_connection, name)

    def initialize(self, connection):
        super().initialize(connection)

        # PGDialect.initialize() checks server version for <= 8.2 and sets
        # this flag to False if so
        if not self.full_returning:
            self.insert_executemany_returning = False

        # HSTORE can't be registered until we have a connection so that
        # we can look up its OID, so we set up this adapter in
        # initialize()
        if self.use_native_hstore:
            info = self._type_info_fetch(connection, "hstore")
            self._has_native_hstore = info is not None
            if self._has_native_hstore:
                from psycopg.types.hstore import register_hstore

                # register the adapter for connections made subsequent to
                # this one
                register_hstore(info, self._psycopg_adapters_map)

                # register the adapter for this connection
                register_hstore(info, connection.connection)

    @classmethod
    def dbapi(cls):
        import psycopg

        return psycopg

    @classmethod
    def get_async_dialect_cls(cls, url):
        return PGDialectAsync_psycopg

    @util.memoized_property
    def _isolation_lookup(self):
        return {
            "READ COMMITTED": self.dbapi.IsolationLevel.READ_COMMITTED,
            "READ UNCOMMITTED": self.dbapi.IsolationLevel.READ_UNCOMMITTED,
            "REPEATABLE READ": self.dbapi.IsolationLevel.REPEATABLE_READ,
            "SERIALIZABLE": self.dbapi.IsolationLevel.SERIALIZABLE,
        }

    @util.memoized_property
    def _psycopg_Json(self):
        from psycopg.types import json

        return json.Json

    @util.memoized_property
    def _psycopg_Jsonb(self):
        from psycopg.types import json

        return json.Jsonb

    @util.memoized_property
    def _psycopg_TransactionStatus(self):
        from psycopg.pq import TransactionStatus

        return TransactionStatus

    @util.memoized_property
    def _psycopg_Range(self):
        from psycopg.types.range import Range

        return Range

    @util.memoized_property
    def _psycopg_Multirange(self):
        from psycopg.types.multirange import Multirange

        return Multirange

    def _do_isolation_level(self, connection, autocommit, isolation_level):
        connection.autocommit = autocommit
        connection.isolation_level = isolation_level

    def get_isolation_level(self, dbapi_connection):
        status_before = dbapi_connection.info.transaction_status
        value = super().get_isolation_level(dbapi_connection)

        # don't rely on psycopg providing enum symbols, compare with
        # eq/ne
        if status_before == self._psycopg_TransactionStatus.IDLE:
            dbapi_connection.rollback()
        return value

    def set_isolation_level(self, dbapi_connection, level):
        if level == "AUTOCOMMIT":
            self._do_isolation_level(dbapi_connection, autocommit=True, isolation_level=None)
        else:
            self._do_isolation_level(
                dbapi_connection,
                autocommit=False,
                isolation_level=self._isolation_lookup[level],
            )

    def set_readonly(self, connection, value):
        connection.read_only = value

    def get_readonly(self, connection):
        return connection.read_only

    def on_connect(self):
        def notices(conn):
            conn.add_notice_handler(_log_notices)

        fns = [notices]

        if self.isolation_level is not None:

            def on_connect(conn):
                self.set_isolation_level(conn, self.isolation_level)

            fns.append(on_connect)

        # fns always has the notices function
        def on_connect(conn):
            for fn in fns:
                fn(conn)

        return on_connect

    def is_disconnect(self, e, connection, cursor):
        if isinstance(e, self.dbapi.Error) and connection is not None:
            if connection.closed or connection.broken:
                return True
        return False

    def _do_prepared_twophase(self, connection, command, recover=False):
        dbapi_conn = connection.connection.dbapi_connection
        if (
            recover
            # don't rely on psycopg providing enum symbols, compare with
            # eq/ne
            or dbapi_conn.info.transaction_status != self._psycopg_TransactionStatus.IDLE
        ):
            dbapi_conn.rollback()
        before_autocommit = dbapi_conn.autocommit
        try:
            if not before_autocommit:
                self._do_autocommit(dbapi_conn, True)
            dbapi_conn.execute(command)
        finally:
            if not before_autocommit:
                self._do_autocommit(dbapi_conn, before_autocommit)

    def do_rollback_twophase(self, connection, xid, is_prepared=True, recover=False):
        if is_prepared:
            self._do_prepared_twophase(connection, f"ROLLBACK PREPARED '{xid}'", recover=recover)
        else:
            self.do_rollback(connection.connection)

    def do_commit_twophase(self, connection, xid, is_prepared=True, recover=False):
        if is_prepared:
            self._do_prepared_twophase(connection, f"COMMIT PREPARED '{xid}'", recover=recover)
        else:
            self.do_commit(connection.connection)

    @util.memoized_property
    def _dialect_specific_select_one(self):
        return ";"


class AsyncAdapt_psycopg_cursor:
    __slots__ = ("_cursor", "await_", "_rows")

    _psycopg_ExecStatus = None

    def __init__(self, cursor, await_) -> None:
        self._cursor = cursor
        self.await_ = await_
        self._rows = []

    def __getattr__(self, name):
        return getattr(self._cursor, name)

    @property
    def arraysize(self):
        return self._cursor.arraysize

    @arraysize.setter
    def arraysize(self, value):
        self._cursor.arraysize = value

    def close(self):
        self._rows.clear()
        # Normal cursor just call _close() in a non-sync way.
        self._cursor._close()

    def execute(self, query, params=None, **kw):
        result = self.await_(self._cursor.execute(query, params, **kw))
        # sqlalchemy result is not async, so need to pull all rows here
        res = self._cursor.pgresult

        # don't rely on psycopg providing enum symbols, compare with
        # eq/ne
        if res and res.status == self._psycopg_ExecStatus.TUPLES_OK:
            rows = self.await_(self._cursor.fetchall())
            if not isinstance(rows, list):
                self._rows = list(rows)
            else:
                self._rows = rows
        return result

    def executemany(self, query, params_seq):
        return self.await_(self._cursor.executemany(query, params_seq))

    def __iter__(self):
        # TODO: try to avoid pop(0) on a list
        while self._rows:
            yield self._rows.pop(0)

    def fetchone(self):
        if self._rows:
            # TODO: try to avoid pop(0) on a list
            return self._rows.pop(0)
        else:
            return None

    def fetchmany(self, size=None):
        if size is None:
            size = self._cursor.arraysize

        retval = self._rows[0:size]
        self._rows = self._rows[size:]
        return retval

    def fetchall(self):
        retval = self._rows
        self._rows = []
        return retval


class AsyncAdapt_psycopg_ss_cursor(AsyncAdapt_psycopg_cursor):
    def execute(self, query, params=None, **kw):
        self.await_(self._cursor.execute(query, params, **kw))
        return self

    def close(self):
        self.await_(self._cursor.close())

    def fetchone(self):
        return self.await_(self._cursor.fetchone())

    def fetchmany(self, size=0):
        return self.await_(self._cursor.fetchmany(size))

    def fetchall(self):
        return self.await_(self._cursor.fetchall())

    def __iter__(self):
        iterator = self._cursor.__aiter__()
        while True:
            try:
                yield self.await_(iterator.__anext__())
            except StopAsyncIteration:
                break


class AsyncAdapt_psycopg_connection(AdaptedConnection):
    __slots__ = ()
    await_ = staticmethod(await_only)

    def __init__(self, connection) -> None:
        self._connection = connection

    def __getattr__(self, name):
        return getattr(self._connection, name)

    def execute(self, query, params=None, **kw):
        cursor = self.await_(self._connection.execute(query, params, **kw))
        return AsyncAdapt_psycopg_cursor(cursor, self.await_)

    def cursor(self, *args, **kw):
        cursor = self._connection.cursor(*args, **kw)
        if hasattr(cursor, "name"):
            return AsyncAdapt_psycopg_ss_cursor(cursor, self.await_)
        else:
            return AsyncAdapt_psycopg_cursor(cursor, self.await_)

    def commit(self):
        self.await_(self._connection.commit())

    def rollback(self):
        self.await_(self._connection.rollback())

    def close(self):
        self.await_(self._connection.close())

    @property
    def autocommit(self):
        return self._connection.autocommit

    @autocommit.setter
    def autocommit(self, value):
        self.set_autocommit(value)

    def set_autocommit(self, value):
        self.await_(self._connection.set_autocommit(value))

    def set_isolation_level(self, value):
        self.await_(self._connection.set_isolation_level(value))

    def set_read_only(self, value):
        self.await_(self._connection.set_read_only(value))

    def set_deferrable(self, value):
        self.await_(self._connection.set_deferrable(value))


class AsyncAdaptFallback_psycopg_connection(AsyncAdapt_psycopg_connection):
    __slots__ = ()
    await_ = staticmethod(await_fallback)


class PsycopgAdaptDBAPI:
    def __init__(self, psycopg) -> None:
        self.psycopg = psycopg

        for k, v in self.psycopg.__dict__.items():
            if k != "connect":
                self.__dict__[k] = v

    def connect(self, *arg, **kw):
        async_fallback = kw.pop("async_fallback", False)
        if util.asbool(async_fallback):
            return AsyncAdaptFallback_psycopg_connection(
                await_fallback(self.psycopg.AsyncConnection.connect(*arg, **kw))
            )
        else:
            return AsyncAdapt_psycopg_connection(await_only(self.psycopg.AsyncConnection.connect(*arg, **kw)))


class PGDialectAsync_psycopg(PGDialect_psycopg):
    is_async = True
    supports_statement_cache = True

    @classmethod
    def dbapi(cls):
        import psycopg
        from psycopg.pq import ExecStatus

        AsyncAdapt_psycopg_cursor._psycopg_ExecStatus = ExecStatus

        return PsycopgAdaptDBAPI(psycopg)

    @classmethod
    def get_pool_class(cls, url):

        async_fallback = url.query.get("async_fallback", False)

        if util.asbool(async_fallback):
            return pool.FallbackAsyncAdaptedQueuePool
        else:
            return pool.AsyncAdaptedQueuePool

    def _type_info_fetch(self, connection, name):
        from psycopg.types import TypeInfo

        adapted = connection.connection
        return adapted.await_(TypeInfo.fetch(adapted.driver_connection, name))

    def _do_isolation_level(self, connection, autocommit, isolation_level):
        connection.set_autocommit(autocommit)
        connection.set_isolation_level(isolation_level)

    def _do_autocommit(self, connection, value):
        connection.set_autocommit(value)

    def set_readonly(self, connection, value):
        connection.set_read_only(value)

    def set_deferrable(self, connection, value):
        connection.set_deferrable(value)

    def get_driver_connection(self, connection):
        return connection._connection


dialect = PGDialect_psycopg
dialect_async = PGDialectAsync_psycopg
