"""Miscellaneous goodies for psycopg2

This module is a generic place used to hold little helper functions
and classes until a better place in the distribution is found.
"""
# psycopg/extras.py - miscellaneous extra goodies for psycopg
#
# Copyright (C) 2003-2019 Federico Di Gregorio  <fog@debian.org>
#
# psycopg2 is free software: you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# In addition, as a special exception, the copyright holders give
# permission to link this program with the OpenSSL library (or with
# modified versions of OpenSSL that use the same license as OpenSSL),
# and distribute linked combinations including the two.
#
# You must obey the GNU Lesser General Public License in all respects for
# all of the code used other than OpenSSL.
#
# psycopg2 is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
# License for more details.

import os as _os
import time as _time
import re as _re
from collections import namedtuple, OrderedDict

import logging as _logging

import psycopg2
from psycopg2 import extensions as _ext
from .extensions import cursor as _cursor
from .extensions import connection as _connection
from .extensions import adapt as _A, quote_ident
from .compat import PY2, PY3, lru_cache

from psycopg2._psycopg import (                             # noqa
    REPLICATION_PHYSICAL, REPLICATION_LOGICAL,
    ReplicationConnection as _replicationConnection,
    ReplicationCursor as _replicationCursor,
    ReplicationMessage)


# expose the json adaptation stuff into the module
from psycopg2._json import (                                # noqa
    json, Json, register_json, register_default_json, register_default_jsonb)


# Expose range-related objects
from psycopg2._range import (                               # noqa
    Range, NumericRange, DateRange, DateTimeRange, DateTimeTZRange,
    register_range, RangeAdapter, RangeCaster)


# Expose ipaddress-related objects
from psycopg2._ipaddress import register_ipaddress          # noqa


class DictCursorBase(_cursor):
    """Base class for all dict-like cursors."""

    def __init__(self, *args, **kwargs):
        if 'row_factory' in kwargs:
            row_factory = kwargs['row_factory']
            del kwargs['row_factory']
        else:
            raise NotImplementedError(
                "DictCursorBase can't be instantiated without a row factory.")
        super(DictCursorBase, self).__init__(*args, **kwargs)
        self._query_executed = False
        self._prefetch = False
        self.row_factory = row_factory

    def fetchone(self):
        if self._prefetch:
            res = super(DictCursorBase, self).fetchone()
        if self._query_executed:
            self._build_index()
        if not self._prefetch:
            res = super(DictCursorBase, self).fetchone()
        return res

    def fetchmany(self, size=None):
        if self._prefetch:
            res = super(DictCursorBase, self).fetchmany(size)
        if self._query_executed:
            self._build_index()
        if not self._prefetch:
            res = super(DictCursorBase, self).fetchmany(size)
        return res

    def fetchall(self):
        if self._prefetch:
            res = super(DictCursorBase, self).fetchall()
        if self._query_executed:
            self._build_index()
        if not self._prefetch:
            res = super(DictCursorBase, self).fetchall()
        return res

    def __iter__(self):
        try:
            if self._prefetch:
                res = super(DictCursorBase, self).__iter__()
                first = next(res)
            if self._query_executed:
                self._build_index()
            if not self._prefetch:
                res = super(DictCursorBase, self).__iter__()
                first = next(res)

            yield first
            while True:
                yield next(res)
        except StopIteration:
            return


class DictConnection(_connection):
    """A connection that uses `DictCursor` automatically."""
    def cursor(self, *args, **kwargs):
        kwargs.setdefault('cursor_factory', DictCursor)
        return super(DictConnection, self).cursor(*args, **kwargs)


class DictCursor(DictCursorBase):
    """A cursor that keeps a list of column name -> index mappings."""

    def __init__(self, *args, **kwargs):
        kwargs['row_factory'] = DictRow
        super(DictCursor, self).__init__(*args, **kwargs)
        self._prefetch = True

    def execute(self, query, vars=None):
        self.index = OrderedDict()
        self._query_executed = True
        return super(DictCursor, self).execute(query, vars)

    def callproc(self, procname, vars=None):
        self.index = OrderedDict()
        self._query_executed = True
        return super(DictCursor, self).callproc(procname, vars)

    def _build_index(self):
        if self._query_executed and self.description:
            for i in range(len(self.description)):
                self.index[self.description[i][0]] = i
            self._query_executed = False


class DictRow(list):
    """A row object that allow by-column-name access to data."""

    __slots__ = ('_index',)

    def __init__(self, cursor):
        self._index = cursor.index
        self[:] = [None] * len(cursor.description)

    def __getitem__(self, x):
        if not isinstance(x, (int, slice)):
            x = self._index[x]
        return super(DictRow, self).__getitem__(x)

    def __setitem__(self, x, v):
        if not isinstance(x, (int, slice)):
            x = self._index[x]
        super(DictRow, self).__setitem__(x, v)

    def items(self):
        g = super(DictRow, self).__getitem__
        return ((n, g(self._index[n])) for n in self._index)

    def keys(self):
        return iter(self._index)

    def values(self):
        g = super(DictRow, self).__getitem__
        return (g(self._index[n]) for n in self._index)

    def get(self, x, default=None):
        try:
            return self[x]
        except Exception:
            return default

    def copy(self):
        return OrderedDict(self.items())

    def __contains__(self, x):
        return x in self._index

    def __getstate__(self):
        return self[:], self._index.copy()

    def __setstate__(self, data):
        self[:] = data[0]
        self._index = data[1]

    if PY2:
        iterkeys = keys
        itervalues = values
        iteritems = items
        has_key = __contains__

        def keys(self):
            return list(self.iterkeys())

        def values(self):
            return tuple(self.itervalues())

        def items(self):
            return list(self.iteritems())


class RealDictConnection(_connection):
    """A connection that uses `RealDictCursor` automatically."""
    def cursor(self, *args, **kwargs):
        kwargs.setdefault('cursor_factory', RealDictCursor)
        return super(RealDictConnection, self).cursor(*args, **kwargs)


class RealDictCursor(DictCursorBase):
    """A cursor that uses a real dict as the base type for rows.

    Note that this cursor is extremely specialized and does not allow
    the normal access (using integer indices) to fetched data. If you need
    to access database rows both as a dictionary and a list, then use
    the generic `DictCursor` instead of `!RealDictCursor`.
    """
    def __init__(self, *args, **kwargs):
        kwargs['row_factory'] = RealDictRow
        super(RealDictCursor, self).__init__(*args, **kwargs)

    def execute(self, query, vars=None):
        self.column_mapping = []
        self._query_executed = True
        return super(RealDictCursor, self).execute(query, vars)

    def callproc(self, procname, vars=None):
        self.column_mapping = []
        self._query_executed = True
        return super(RealDictCursor, self).callproc(procname, vars)

    def _build_index(self):
        if self._query_executed and self.description:
            self.column_mapping = [d[0] for d in self.description]
            self._query_executed = False


class RealDictRow(OrderedDict):
    """A `!dict` subclass representing a data record."""

    def __init__(self, *args, **kwargs):
        if args and isinstance(args[0], _cursor):
            cursor = args[0]
            args = args[1:]
        else:
            cursor = None

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

        if cursor is not None:
            # Required for named cursors
            if cursor.description and not cursor.column_mapping:
                cursor._build_index()

            # Store the cols mapping in the dict itself until the row is fully
            # populated, so we don't need to add attributes to the class
            # (hence keeping its maintenance, special pickle support, etc.)
            self[RealDictRow] = cursor.column_mapping

    def __setitem__(self, key, value):
        if RealDictRow in self:
            # We are in the row building phase
            mapping = self[RealDictRow]
            super(RealDictRow, self).__setitem__(mapping[key], value)
            if key == len(mapping) - 1:
                # Row building finished
                del self[RealDictRow]
            return

        super(RealDictRow, self).__setitem__(key, value)


class NamedTupleConnection(_connection):
    """A connection that uses `NamedTupleCursor` automatically."""
    def cursor(self, *args, **kwargs):
        kwargs.setdefault('cursor_factory', NamedTupleCursor)
        return super(NamedTupleConnection, self).cursor(*args, **kwargs)


class NamedTupleCursor(_cursor):
    """A cursor that generates results as `~collections.namedtuple`.

    `!fetch*()` methods will return named tuples instead of regular tuples, so
    their elements can be accessed both as regular numeric items as well as
    attributes.

        >>> nt_cur = conn.cursor(cursor_factory=psycopg2.extras.NamedTupleCursor)
        >>> rec = nt_cur.fetchone()
        >>> rec
        Record(id=1, num=100, data="abc'def")
        >>> rec[1]
        100
        >>> rec.data
        "abc'def"
    """
    Record = None
    MAX_CACHE = 1024

    def execute(self, query, vars=None):
        self.Record = None
        return super(NamedTupleCursor, self).execute(query, vars)

    def executemany(self, query, vars):
        self.Record = None
        return super(NamedTupleCursor, self).executemany(query, vars)

    def callproc(self, procname, vars=None):
        self.Record = None
        return super(NamedTupleCursor, self).callproc(procname, vars)

    def fetchone(self):
        t = super(NamedTupleCursor, self).fetchone()
        if t is not None:
            nt = self.Record
            if nt is None:
                nt = self.Record = self._make_nt()
            return nt._make(t)

    def fetchmany(self, size=None):
        ts = super(NamedTupleCursor, self).fetchmany(size)
        nt = self.Record
        if nt is None:
            nt = self.Record = self._make_nt()
        return list(map(nt._make, ts))

    def fetchall(self):
        ts = super(NamedTupleCursor, self).fetchall()