'''
Created on Jun 12, 2020

@author Saurabh Kumar Pandey
'''

from pymongo import MongoClient,errors

#Local application/Library specific imports 
import config as cfg

class Pool:
    def __init__(self, address, options, handshake=True):
        """
        :Parameters:
          - `address`: a (hostname, port) tuple
          - `options`: a PoolOptions instance
          - `handshake`: whether to call ismaster for each new SocketInfo
        """
        # Check a socket's health with socket_closed() every once in a while.
        # Can override for testing: 0 to always check, None to never check.
        self._check_interval_seconds = 1
        # LIFO pool. Sockets are ordered on idle time. Sockets claimed
        # and returned to pool from the left side. Stale sockets removed
        # from the right side.
        self.sockets = collections.deque()
        self.lock = threading.Lock()
        self.active_sockets = 0
        # Monotonically increasing connection ID required for CMAP Events.
        self.next_connection_id = 1
        self.closed = False
        # Track whether the sockets in this pool are writeable or not.
        self.is_writable = None

        # Keep track of resets, so we notice sockets created before the most
        # recent reset and close them.
        self.generation = 0
        self.pid = os.getpid()
        self.address = address
        self.opts = options
        self.handshake = handshake
        # Don't publish events in Monitor pools.
        self.enabled_for_cmap = (
                self.handshake and
                self.opts.event_listeners is not None and
                self.opts.event_listeners.enabled_for_cmap)

        if (self.opts.wait_queue_multiple is None or
                self.opts.max_pool_size is None):
            max_waiters = None
        else:
            max_waiters = (
                self.opts.max_pool_size * self.opts.wait_queue_multiple)

        self._socket_semaphore = thread_util.create_semaphore(
            self.opts.max_pool_size, max_waiters)
        if self.enabled_for_cmap:
            self.opts.event_listeners.publish_pool_created(
                self.address, self.opts.non_default_options)

    def _reset(self, close):
        with self.lock:
            if self.closed:
                return
            self.generation += 1
            self.pid = os.getpid()
            sockets, self.sockets = self.sockets, collections.deque()
            self.active_sockets = 0
            if close:
                self.closed = True

        listeners = self.opts.event_listeners
        # CMAP spec says that close() MUST close sockets before publishing the
        # PoolClosedEvent but that reset() SHOULD close sockets *after*
        # publishing the PoolClearedEvent.
        if close:
            for sock_info in sockets:
                sock_info.close_socket(ConnectionClosedReason.POOL_CLOSED)
            if self.enabled_for_cmap:
                listeners.publish_pool_closed(self.address)
        else:
            if self.enabled_for_cmap:
                listeners.publish_pool_cleared(self.address)
            for sock_info in sockets:
                sock_info.close_socket(ConnectionClosedReason.STALE)

    def update_is_writable(self, is_writable):
        """Updates the is_writable attribute on all sockets currently in the
        Pool.
        """
        self.is_writable = is_writable
        with self.lock:
            for socket in self.sockets:
                socket.update_is_writable(self.is_writable)

    def reset(self):
        self._reset(close=False)

    def close(self):
        self._reset(close=True)

    def remove_stale_sockets(self, reference_generation, all_credentials):
        """Removes stale sockets then adds new ones if pool is too small and
        has not been reset. The `reference_generation` argument specifies the
        `generation` at the point in time this operation was requested on the
        pool.
        """
        if self.opts.max_idle_time_seconds is not None:
            with self.lock:
                while (self.sockets and
                       self.sockets[-1].idle_time_seconds() > self.opts.max_idle_time_seconds):
                    sock_info = self.sockets.pop()
                    sock_info.close_socket(ConnectionClosedReason.IDLE)

        while True:
            with self.lock:
                if (len(self.sockets) + self.active_sockets >=
                        self.opts.min_pool_size):
                    # There are enough sockets in the pool.
                    break

            # We must acquire the semaphore to respect max_pool_size.
            if not self._socket_semaphore.acquire(False):
                break
            try:
                sock_info = self.connect(all_credentials)
                with self.lock:
                    # Close connection and return if the pool was reset during
                    # socket creation or while acquiring the pool lock.
                    if self.generation != reference_generation:
                        sock_info.close_socket(ConnectionClosedReason.STALE)
                        break
                    self.sockets.appendleft(sock_info)
            finally:
                self._socket_semaphore.release()

    def connect(self, all_credentials=None):
        """Connect to Mongo and return a new SocketInfo.
        Can raise ConnectionFailure or CertificateError.
        Note that the pool does not keep a reference to the socket -- you
        must call return_socket() when you're done with it.
        """
        with self.lock:
            conn_id = self.next_connection_id
            self.next_connection_id += 1

        listeners = self.opts.event_listeners
        if self.enabled_for_cmap:
            listeners.publish_connection_created(self.address, conn_id)

        sock = None
        try:
            sock = _configured_socket(self.address, self.opts)
        except socket.error as error:
            if self.enabled_for_cmap:
                listeners.publish_connection_closed(
                    self.address, conn_id, ConnectionClosedReason.ERROR)

            _raise_connection_failure(self.address, error)

        sock_info = SocketInfo(sock, self, self.address, conn_id)
        if self.handshake:
            sock_info.ismaster(self.opts.metadata, None, all_credentials)
            self.is_writable = sock_info.is_writable

        return sock_info

    def get_socket(self, all_credentials, checkout=False):
        """Get a socket from the pool. Use with a "with" statement.
        Returns a :class:`SocketInfo` object wrapping a connected
        :class:`socket.socket`.
        This method should always be used in a with-statement::
            with pool.get_socket(credentials, checkout) as socket_info:
                socket_info.send_message(msg)
                data = socket_info.receive_message(op_code, request_id)
        The socket is logged in or out as needed to match ``all_credentials``
        using the correct authentication mechanism for the server's wire
        protocol version.
        Can raise ConnectionFailure or OperationFailure.
        :Parameters:
          - `all_credentials`: dict, maps auth source to MongoCredential.
          - `checkout` (optional): keep socket checked out.
        """
        listeners = self.opts.event_listeners
        if self.enabled_for_cmap:
            listeners.publish_connection_check_out_started(self.address)

        sock_info = self._get_socket(all_credentials)

        if self.enabled_for_cmap:
            listeners.publish_connection_checked_out(
                self.address, sock_info.id)
        try:
            yield sock_info
        except:
            # Exception in caller. Decrement semaphore.
            self.return_socket(sock_info)
            raise
        else:
            if not checkout:
                self.return_socket(sock_info)

    def _get_socket(self, all_credentials):
        """Get or create a SocketInfo. Can raise ConnectionFailure."""
        # We use the pid here to avoid issues with fork / multiprocessing.
        # See test.test_client:TestClient.test_fork for an example of
        # what could go wrong otherwise
        if self.pid != os.getpid():
            self.reset()

        if self.closed:
            if self.enabled_for_cmap:
                self.opts.event_listeners.publish_connection_check_out_failed(
                    self.address, ConnectionCheckOutFailedReason.POOL_CLOSED)
            raise _PoolClosedError(
                'Attempted to check out a connection from closed connection '
                'pool')

        # Get a free socket or create one.
        if not self._socket_semaphore.acquire(
                True, self.opts.wait_queue_timeout):
            self._raise_wait_queue_timeout()
        with self.lock:
            self.active_sockets += 1

        # We've now acquired the semaphore and must release it on error.
        try:
            sock_info = None
            while sock_info is None:
                try:
                    with self.lock:
                        sock_info = self.sockets.popleft()
                except IndexError:
                    # Can raise ConnectionFailure or CertificateError.
                    sock_info = self.connect(all_credentials)
                else:
                    if self._perished(sock_info):
                        sock_info = None
            sock_info.check_auth(all_credentials)
        except Exception:
            self._socket_semaphore.release()
            with self.lock:
                self.active_sockets -= 1

            if self.enabled_for_cmap:
                self.opts.event_listeners.publish_connection_check_out_failed(
                    self.address, ConnectionCheckOutFailedReason.CONN_ERROR)
            raise

        return sock_info

    def return_socket(self, sock_info):
        """Return the socket to the pool, or if it's closed discard it.
        :Parameters:
          - `sock_info`: The socket to check into the pool.
        """
        listeners = self.opts.event_listeners
        if self.enabled_for_cmap:
            listeners.publish_connection_checked_in(self.address, sock_info.id)
        if self.pid != os.getpid():
            self.reset()
        else:
            if self.closed:
                sock_info.close_socket(ConnectionClosedReason.POOL_CLOSED)
            elif not sock_info.closed:
                with self.lock:
                    # Hold the lock to ensure this section does not race with
                    # Pool.reset().
                    if sock_info.generation != self.generation:
                        sock_info.close_socket(ConnectionClosedReason.STALE)
                    else:
                        sock_info.update_last_checkin_time()
                        sock_info.update_is_writable(self.is_writable)
                        self.sockets.appendleft(sock_info)

        self._socket_semaphore.release()
        with self.lock:
            self.active_sockets -= 1

    def _perished(self, sock_info):
        """Return True and close the connection if it is "perished".
        This side-effecty function checks if this socket has been idle for
        for longer than the max idle time, or if the socket has been closed by
        some external network error, or if the socket's generation is outdated.
        Checking sockets lets us avoid seeing *some*
        :class:`~pymongo.errors.AutoReconnect` exceptions on server
        hiccups, etc. We only check if the socket was closed by an external
        error if it has been > 1 second since the socket was checked into the
        pool, to keep performance reasonable - we can't avoid AutoReconnects
        completely anyway.
        """
        idle_time_seconds = sock_info.idle_time_seconds()
        # If socket is idle, open a new one.
        if (self.opts.max_idle_time_seconds is not None and
                idle_time_seconds > self.opts.max_idle_time_seconds):
            sock_info.close_socket(ConnectionClosedReason.IDLE)
            return True

        if (self._check_interval_seconds is not None and (
                0 == self._check_interval_seconds or
                idle_time_seconds > self._check_interval_seconds)):
            if sock_info.socket_closed():
                sock_info.close_socket(ConnectionClosedReason.ERROR)
                return True

        if sock_info.generation != self.generation:
            sock_info.close_socket(ConnectionClosedReason.STALE)
            return True

        return False

    def _raise_wait_queue_timeout(self):
        listeners = self.opts.event_listeners
        if self.enabled_for_cmap:
            listeners.publish_connection_check_out_failed(
                self.address, ConnectionCheckOutFailedReason.TIMEOUT)
        raise ConnectionFailure(
            'Timed out while checking out a connection from connection pool '
            'with max_size %r and wait_queue_timeout %r' % (
                self.opts.max_pool_size, self.opts.wait_queue_timeout))

    def __del__(self):
        # Avoid ResourceWarnings in Python 3
        # Close all sockets without calling reset() or close() because it is
        # not safe to acquire a lock in __del__.
        for sock_info in self.sockets:
            sock_info.close_socket(None)

mongolib = Pool()
mongolib.connect()            
