"""
Connection Pooling

"""
import datetime
import logging
import os
import threading
import time
import weakref

import psycopg2

LOGGER = logging.getLogger(__name__)

DEFAULT_IDLE_TTL = 60
DEFAULT_MAX_SIZE = int(os.environ.get('QUERIES_MAX_POOL_SIZE', 1))


class Connection(object):
    """Contains the handle to the connection, the current state of the
    connection and methods for manipulating the state of the connection.

    """
    _lock = threading.Lock()

    def __init__(self, handle):
        self.handle = handle
        self.used_by = None
        self.executions = 0
        self.exceptions = 0

    def close(self):
        """Close the connection

        :raises: ConnectionBusyError

        """
        LOGGER.debug('Connection %s closing', self.id)
        if self.busy and not self.closed:
            raise ConnectionBusyError(self)
        with self._lock:
            if not self.handle.closed:
                try:
                    self.handle.close()
                except psycopg2.InterfaceError as error:
                    LOGGER.error('Error closing socket: %s', error)

    @property
    def closed(self):
        """Return if the psycopg2 connection is closed.

        :rtype: bool

        """
        return self.handle.closed != 0

    @property
    def busy(self):
        """Return if the connection is currently executing a query or is locked
        by a session that still exists.

        :rtype: bool

        """
        if self.handle.isexecuting():
            return True
        elif self.used_by is None:
            return False
        return not self.used_by() is None

    @property
    def executing(self):
        """Return if the connection is currently executing a query

        :rtype: bool

        """
        return self.handle.isexecuting()

    def free(self):
        """Remove the lock on the connection if the connection is not active

        :raises: ConnectionBusyError

        """
        LOGGER.debug('Connection %s freeing', self.id)
        if self.handle.isexecuting():
            raise ConnectionBusyError(self)
        with self._lock:
            self.used_by = None
        LOGGER.debug('Connection %s freed', self.id)

    @property
    def id(self):
        """Return id of the psycopg2 connection object

        :rtype: int

        """
        return id(self.handle)

    def lock(self, session):
        """Lock the connection, ensuring that it is not busy and storing
        a weakref for the session.

        :param queries.Session session: The session to lock the connection with
        :raises: ConnectionBusyError

        """
        if self.busy:
            raise ConnectionBusyError(self)
        with self._lock:
            self.used_by = weakref.ref(session)
        LOGGER.debug('Connection %s locked', self.id)

    @property
    def locked(self):
        """Return if the connection is currently exclusively locked

        :rtype: bool

        """
        return self.used_by is not None


class Pool(object):
    """A connection pool for gaining access to and managing connections"""
    _lock = threading.Lock()

    idle_start = None
    idle_ttl = DEFAULT_IDLE_TTL
    max_size = DEFAULT_MAX_SIZE

    def __init__(self,
                 pool_id,
                 idle_ttl=DEFAULT_IDLE_TTL,
                 max_size=DEFAULT_MAX_SIZE,
                 time_method=None):
        self.connections = {}
        self._id = pool_id
        self.idle_ttl = idle_ttl
        self.max_size = max_size
        self.time_method = time_method or time.time

    def __contains__(self, connection):
        """Return True if the pool contains the connection"""
        return id(connection) in self.connections

    def __len__(self):
        """Return the number of connections in the pool"""
        return len(self.connections)

    def add(self, connection):
        """Add a new connection to the pool

        :param connection: The connection to add to the pool
        :type connection: psycopg2.extensions.connection
        :raises: PoolFullError

        """
        if id(connection) in self.connections:
            raise ValueError('Connection already exists in pool')

        if len(self.connections) == self.max_size:
            LOGGER.warning('Race condition found when adding new connection')
            try:
                connection.close()
            except (psycopg2.Error, psycopg2.Warning) as error:
                LOGGER.error('Error closing the conn that cant be used: %s',
                             error)
            raise PoolFullError(self)
        with self._lock:
            self.connections[id(connection)] = Connection(connection)
        LOGGER.debug('Pool %s added connection %s', self.id, id(connection))

    @property
    def busy_connections(self):
        """Return a list of active/busy connections

        :rtype: list

        """
        return [c for c in self.connections.values()
                if c.busy and not c.closed]

    def clean(self):
        """Clean the pool by removing any closed connections and if the pool's
        idle has exceeded its idle TTL, remove all connections.

        """
        LOGGER.debug('Cleaning the pool')
        for connection in [self.connections[k] for k in self.connections if
                           self.connections[k].closed]:
            LOGGER.debug('Removing %s', connection.id)
            self.remove(connection.handle)

        if self.idle_duration > self.idle_ttl:
            self.close()

        LOGGER.debug('Pool %s cleaned', self.id)

    def close(self):
        """Close the pool by closing and removing all of the connections"""
        for cid in list(self.connections.keys()):
            self.remove(self.connections[cid].handle)
        LOGGER.debug('Pool %s closed', self.id)

    @property
    def closed_connections(self):
        """Return a list of closed connections

        :rtype: list

        """
        return [c for c in self.connections.values() if c.closed]

    def connection_handle(self, connection):
        """Return a connection object for the given psycopg2 connection

        :param connection: The connection to return a parent for
        :type connection: psycopg2.extensions.connection
        :rtype: Connection

        """
        return self.connections[id(connection)]

    @property
    def executing_connections(self):
        """Return a list of connections actively executing queries

        :rtype: list

        """
        return [c for c in self.connections.values() if c.executing]

    def free(self, connection):
        """Free the connection from use by the session that was using it.

        :param connection: The connection to free
        :type connection: psycopg2.extensions.connection
        :raises: ConnectionNotFoundError

        """
        LOGGER.debug('Pool %s freeing connection %s', self.id, id(connection))
        try:
            self.connection_handle(connection).free()
        except KeyError:
            raise ConnectionNotFoundError(self.id, id(connection))

        if self.idle_connections == list(self.connections.values()):
            with self._lock:
                self.idle_start = self.time_method()
        LOGGER.debug('Pool %s freed connection %s', self.id, id(connection))

    def get(self, session):
        """Return an idle connection and assign the session to the connection

        :param queries.Session session: The session to assign
        :rtype: psycopg2.extensions.connection
        :raises: NoIdleConnectionsError

        """
        idle = self.idle_connections
        if idle:
            connection = idle.pop(0)
            connection.lock(session)
            if self.idle_start:
                with self._lock:
                    self.idle_start = None
            return connection.handle
        raise NoIdleConnectionsError(self.id)

    @property
    def id(self):
        """Return the ID for this pool

        :rtype: str

        """
        return self._id

    @property
    def idle_connections(self):
        """Return a list of idle connections

        :rtype: list

        """
        return [c for c in self.connections.values()
                if not c.busy and not c.closed]

    @property
    def idle_duration(self):
        """Return the number of seconds that the pool has had no active
        connections.

        :rtype: float

        """
        if self.idle_start is None:
            return 0
        return self.time_method() - self.idle_start

    @property
    def is_full(self):
        """Return True if there are no more open slots for connections.

        :rtype: bool

        """
        return len(self.connections) >= self.max_size

    def lock(self, connection, session):
        """Explicitly lock the specified connection

        :type connection: psycopg2.extensions.connection
        :param connection: The connection to lock
        :param queries.Session session: The session to hold the lock

        """
        cid = id(connection)
        try:
            self.connection_handle(connection).lock(session)
        except KeyError:
            raise ConnectionNotFoundError(self.id, cid)
        else:
            if self.idle_start:
                with self._lock:
                    self.idle_start = None
        LOGGER.debug('Pool %s locked connection %s', self.id, cid)

    @property
    def locked_connections(self):
        """Return a list of all locked connections

        :rtype: list

        """
        return [c for c in self.connections.values() if c.locked]

    def remove(self, connection):
        """Remove the connection from the pool

        :param connection: The connection to remove
        :type connection: psycopg2.extensions.connection
        :raises: ConnectionNotFoundError
        :raises: ConnectionBusyError

        """
        cid = id(connection)
        if cid not in self.connections:
            raise ConnectionNotFoundError(self.id, cid)
        self.connection_handle(connection).close()
        with self._lock:
            del self.connections[cid]
        LOGGER.debug('Pool %s removed connection %s', self.id, cid)

    def report(self):
        """Return a report about the pool state and configuration.

        :rtype: dict

        """
        return {
            'connections': {
                'busy': len(self.busy_connections),
                'closed': len(self.closed_connections),
                'executing': len(self.executing_connections),
                'idle': len(self.idle_connections),
                'locked': len(self.busy_connections)
            },
            'exceptions': sum([c.exceptions
                               for c in self.connections.values()]),
            'executions': sum([c.executions
                               for c in self.connections.values()]),
            'full': self.is_full,
            'idle': {
                'duration': self.idle_duration,
                'ttl': self.idle_ttl
            },
            'max_size': self.max_size
        }

    def shutdown(self):
        """Forcefully shutdown the entire pool, closing all non-executing
        connections.

        :raises: ConnectionBusyError

        """
        with self._lock:
            for cid in list(self.connections.keys()):
                if self.connections[cid].executing:
                    raise ConnectionBusyError(cid)
                if self.connections[cid].locked:
                    self.connections[cid].free()
                self.connections[cid].close()
                del self.connections[cid]

    def set_idle_ttl(self, ttl):
        """Set the idle ttl

        :param int ttl: The TTL when idle

        """
        with self._lock:
            self.idle_ttl = ttl

    def set_max_size(self, size):
        """Set the maximum number of connections

        :param int size: The maximum number of connections

        """
        with self._lock:
            self.max_size = size


class PoolManager(object):
    """The connection pool object implements behavior around connections and
    their use in queries.Session objects.

    We carry a pool id instead of the connection URI so that we will not be
    carrying the URI in memory, creating a possible security issue.

    """
    _lock = threading.Lock()
    _pools = {}

    def __contains__(self, pid):
        """Returns True if the pool exists

        :param str pid: The pool id to check for
        :rtype: bool

        """
        return pid in self.__class__._pools

    @classmethod
    def instance(cls):
        """Only allow a single PoolManager instance to exist, returning the
        handle for it.

        :rtype: PoolManager

        """
        if not hasattr(cls, '_instance'):
            with cls._lock:
                cls._instance = cls()
        return cls._instance

    @classmethod
    def add(cls, pid, connection):
        """Add a new connection and session to a pool.

        :param str pid: The pool id
        :type connection: psycopg2.extensions.connection
        :param connection: The connection to add to the pool

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            cls._pools[pid].add(connection)

    @classmethod
    def clean(cls, pid):
        """Clean the specified pool, removing any closed connections or
        stale locks.

        :param str pid: The pool id to clean

        """
        with cls._lock:
            try:
                cls._ensure_pool_exists(pid)
            except KeyError:
                LOGGER.debug('Pool clean invoked against missing pool %s', pid)
                return
            cls._pools[pid].clean()
            cls._maybe_remove_pool(pid)

    @classmethod
    def create(cls, pid, idle_ttl=DEFAULT_IDLE_TTL, max_size=DEFAULT_MAX_SIZE,
               time_method=None):
        """Create a new pool, with the ability to pass in values to override
        the default idle TTL and the default maximum size.

        A pool's idle TTL defines the amount of time that a pool can be open
        without any sessions before it is removed.

        A pool's max size defines the maximum number of connections that can
        be added to the pool to prevent unbounded open connections.

        :param str pid: The pool ID
        :param int idle_ttl: Time in seconds for the idle TTL
        :param int max_size: The maximum pool size
        :param callable time_method: Override the use of :py:meth:`time.time`
            method for time values.
        :raises: KeyError

        """
        if pid in cls._pools:
            raise KeyError('Pool %s already exists' % pid)
        with cls._lock:
            LOGGER.debug("Creating Pool: %s (%i/%i)", pid, idle_ttl, max_size)
            cls._pools[pid] = Pool(pid, idle_ttl, max_size, time_method)

    @classmethod
    def free(cls, pid, connection):
        """Free a connection that was locked by a session

        :param str pid: The pool ID
        :param connection: The connection to remove
        :type connection: psycopg2.extensions.connection

        """
        with cls._lock:
            LOGGER.debug('Freeing %s from pool %s', id(connection), pid)
            cls._ensure_pool_exists(pid)
            cls._pools[pid].free(connection)

    @classmethod
    def get(cls, pid, session):
        """Get an idle, unused connection from the pool. Once a connection has
        been retrieved, it will be marked as in-use until it is freed.

        :param str pid: The pool ID
        :param queries.Session session: The session to assign to the connection
        :rtype: psycopg2.extensions.connection

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            return cls._pools[pid].get(session)

    @classmethod
    def get_connection(cls, pid, connection):
        """Return the specified :class:`~queries.pool.Connection` from the
        pool.

        :param str pid: The pool ID
        :param connection: The connection to return for
        :type connection: psycopg2.extensions.connection
        :rtype: queries.pool.Connection

        """
        with cls._lock:
            return cls._pools[pid].connection_handle(connection)

    @classmethod
    def has_connection(cls, pid, connection):
        """Check to see if a pool has the specified connection

        :param str pid: The pool ID
        :param connection: The connection to check for
        :type connection: psycopg2.extensions.connection
        :rtype: bool

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            return connection in cls._pools[pid]

    @classmethod
    def has_idle_connection(cls, pid):
        """Check to see if a pool has an idle connection

        :param str pid: The pool ID
        :rtype: bool

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            return bool(cls._pools[pid].idle_connections)

    @classmethod
    def is_full(cls, pid):
        """Return a bool indicating if the specified pool is full

        :param str pid: The pool id
        :rtype: bool

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            return cls._pools[pid].is_full

    @classmethod
    def lock(cls, pid, connection, session):
        """Explicitly lock the specified connection in the pool

        :param str pid: The pool id
        :type connection: psycopg2.extensions.connection
        :param connection: The connection to add to the pool
        :param queries.Session session: The session to hold the lock

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            cls._pools[pid].lock(connection, session)

    @classmethod
    def remove(cls, pid):
        """Remove a pool, closing all connections

        :param str pid: The pool ID

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            cls._pools[pid].close()
            del cls._pools[pid]

    @classmethod
    def remove_connection(cls, pid, connection):
        """Remove a connection from the pool, closing it if is open.

        :param str pid: The pool ID
        :param connection: The connection to remove
        :type connection: psycopg2.extensions.connection
        :raises: ConnectionNotFoundError

        """
        cls._ensure_pool_exists(pid)
        cls._pools[pid].remove(connection)

    @classmethod
    def set_idle_ttl(cls, pid, ttl):
        """Set the idle TTL for a pool, after which it will be destroyed.

        :param str pid: The pool id
        :param int ttl: The TTL for an idle pool

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            cls._pools[pid].set_idle_ttl(ttl)

    @classmethod
    def set_max_size(cls, pid, size):
        """Set the maximum number of connections for the specified pool

        :param str pid: The pool to set the size for
        :param int size: The maximum number of connections

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            cls._pools[pid].set_max_size(size)

    @classmethod
    def shutdown(cls):
        """Close all connections on in all pools"""
        for pid in list(cls._pools.keys()):
            cls._pools[pid].shutdown()
        LOGGER.info('Shutdown complete, all pooled connections closed')

    @classmethod
    def size(cls, pid):
        """Return the number of connections in the pool

        :param str pid: The pool id
        :rtype int

        """
        with cls._lock:
            cls._ensure_pool_exists(pid)
            return len(cls._pools[pid])

    @classmethod
    def report(cls):
        """Return the state of the all of the registered pools.

        :rtype: dict

        """
        return {
            'timestamp': datetime.datetime.utcnow().isoformat(),
            'process': os.getpid(),
            'pools': dict([(i, p.report()) for i, p in cls._pools.items()])
        }

    @classmethod
    def _ensure_pool_exists(cls, pid):
        """Raise an exception if the pool has yet to be created or has been
        removed.

        :param str pid: The pool ID to check for
        :raises: KeyError

        """
        if pid not in cls._pools:
            raise KeyError('Pool %s has not been created' % pid)

    @classmethod
    def _maybe_remove_pool(cls, pid):
        """If the pool has no open connections, remove it

        :param str pid: The pool id to clean

        """
        if not len(cls._pools[pid]):
            del cls._pools[pid]


class QueriesException(Exception):
    """Base Exception for all other Queries exceptions"""
    pass


class ConnectionException(QueriesException):
    def __init__(self, cid):
        self.cid = cid


class PoolException(QueriesException):
    def __init__(self, pid):
        self.pid = pid


class PoolConnectionException(PoolException):
    def __init__(self, pid, cid):
        self.pid = pid
        self.cid = cid


class ActivePoolError(PoolException):
    """Raised when removing a pool that has active connections"""
    def __str__(self):
        return 'Pool %s has at least one active connection' % self.pid


class ConnectionBusyError(ConnectionException):
    """Raised when trying to lock a connection that is already busy"""
    def __str__(self):
        return 'Connection %s is busy' % self.cid


class ConnectionNotFoundError(PoolConnectionException):
    """Raised if a specific connection is not found in the pool"""
    def __str__(self):
        return 'Connection %s not found in pool %s' % (self.cid, self.pid)


class NoIdleConnectionsError(PoolException):
    """Raised if a pool does not have any idle, open connections"""
    def __str__(self):
        return 'Pool %s has no idle connections' % self.pid


class PoolFullError(PoolException):
    """Raised when adding a connection to a pool that has hit max-size"""
    def __str__(self):
        return 'Pool %s is at its maximum capacity' % self.pid
