from typing import Any, List, Dict

from psycopg2.pool import ThreadedConnectionPool
from .._ThreadPoolConfig import ThreadPoolConfig


class Psycopg2MultipleThreadPool(object):
    def __init__(self, configs):
        # type: (List[ThreadPoolConfig]) -> None
        self.config_dict = {}
        for config in configs:
            self.config_dict[config.key] = config

        self.thread_pool_dict = {}

        self.conn_pool_dict = {}  # type: Dict[Any, ThreadedConnectionPool]

    def get_conn(self, key):
        pool = self.get_thread_pool(key)
        conn = pool.getconn()
        self.conn_pool_dict[conn] = pool

        return conn

    def put_conn(self, conn):
        self.conn_pool_dict[conn].putconn(conn)

    def get_thread_pool(self, key):
        # type: (str) -> ThreadedConnectionPool
        try:
            return self.thread_pool_dict[key]
        except KeyError:
            try:
                config = self.config_dict[key]
                pool = ThreadedConnectionPool(config.min_connections,
                                              config.max_connections,
                                              host=config.host,
                                              port=config.port,
                                              database=config.database,
                                              user=config.username,
                                              password=config.password)
                self.thread_pool_dict[key] = pool
                return pool
            except KeyError:
                raise Exception("The key {} is not configured in this thread pool".format(key))
