import threading
from collections import deque
from contextlib import contextmanager

from funcy import memoize, wrap_with


@contextmanager
def get_connection(conn_func, *args, **kwargs):
    pool = get_pool(conn_func, *args, **kwargs)
    conn = pool.get_connection()
    try:
        yield conn
    except BaseException:
        conn.close()
        raise
    else:
        pool.release(conn)


@wrap_with(threading.Lock())
@memoize
def get_pool(conn_func, *args, **kwargs):
    return Pool(conn_func, *args, **kwargs)


def close_pools():
    for pool in get_pool.memory.values():
        pool.close()
    get_pool.memory.clear()


class Pool:
    def __init__(self, conn_func, *conn_args, **conn_kwargs):
        self._conn_func = conn_func
        self._conn_args = conn_args
        self._conn_kwargs = conn_kwargs
        self._conns = deque()
        self._closed = False

    def __del__(self):
        self.close()

    def close(self):
        while self._conns:
            self._conns.pop().close()
        self._closed = True

    def get_connection(self):
        try:
            return self._conns.popleft()
        except IndexError:
            return self._conn_func(*self._conn_args, **self._conn_kwargs)

    def release(self, conn):
        if self._closed:
            conn.close()
        else:
            self._conns.append(conn)
