import atexit
import threading
import logging
from contextlib import contextmanager

from waddle import ParamBunch
from .db2i import Db2i
from .hp3000 import Hp3000
from .mssql import SqlServer
from .mysql import Mysql
from .netsuite import Netsuite
from .oracle import Oracle
from .postgres import Postgres
from .redshift import Redshift
from .snowflake import Snowflake
from .vertica import Vertica


gtl = threading.local()
log = logging.getLogger(__name__)



class ConnectionManager:
    thread_local_key = 'concentric_connection_manager'


    def __init__(self):
        self.conf = ParamBunch()

    def add_config(self, *filenames):
        for filename in filenames:
            self.conf.from_file(filename, True)

    @classmethod
    def get_instance(cls, *filenames):
        global gtl
        manager = getattr(gtl, cls.thread_local_key, None)
        if manager:
            return manager
        manager = cls()
        manager.add_config(*filenames)
        setattr(gtl, cls.thread_local_key, manager)
        return manager

    @classmethod
    def get_connector(cls, alias):
        manager = cls.get_instance()
        conf = manager.conf[alias]
        if conf.engine == 'oracle':
            return conf, Oracle
        if conf.engine == 'db2i':
            return conf, Db2i
        if conf.engine == 'hp3000':
            return conf, Hp3000
        if conf.engine == 'mssql':
            return conf, SqlServer
        if conf.engine == 'mysql':
            return conf, Mysql
        if conf.engine == 'netsuite':
            return conf, Netsuite
        if conf.engine == 'postgres':
            return conf, Postgres
        if conf.engine == 'redshift':
            return conf, Redshift
        if conf.engine == 'snowflake':
            return conf, Snowflake
        if conf.engine == 'vertica':
            return conf, Vertica
        return None, None

    @classmethod
    def connect(cls, alias, ping=False):
        conf, connector = cls.get_connector(alias)
        return connector.connect(conf)


class CachingConnectionManager(ConnectionManager):
    thread_local_key = 'concentric_caching_connection_manager'
    thread_local_cache_key = 'concentric_connection_cache'

    @classmethod
    def test_connection(cls, alias, conn):
        _, engine = cls.get_connector(alias)
        ping_query = getattr(engine, 'ping_query')
        if ping_query:
            try:
                with conn.cursor() as cursor:
                    cursor.execute(ping_query)
                    cursor.fetchall()
            except:
                conn = None
        return conn

    @classmethod
    def connect(cls, alias, use_cache=True, pre_ping=True):
        global gtl
        cache = getattr(gtl, cls.thread_local_cache_key, None) or {}
        if use_cache:
            conn = cache.get(alias)
            if pre_ping and conn:
                if not hasattr(conn, 'ping'):
                    conn = cls.test_connection(alias, conn)
                else:
                    try:
                        conn.ping()
                    except:
                        conn = None
            if conn:
                return conn
        conn = super().connect(alias)
        if use_cache:
            cache[alias] = conn
            setattr(gtl, cls.thread_local_cache_key, cache)
        return conn

    @classmethod
    def cleanup(cls):
        global gtl
        cache = getattr(gtl, cls.thread_local_cache_key, None) or {}
        keys = []
        for alias, conn in cache.items():
            try:
                keys.append(alias)
                conn.close()
            except:
                pass
        for alias in keys:
            cache.pop(alias, None)


@contextmanager
def transactional(session_cls, initial):
    """
    Provide a transactional scope around a series of operations.
    """
    session = session_cls()
    try:
        for x in initial:
            session.execute(x)
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager
def read_only(session_cls, initial, **params):
    """
    Provide a closing scope around a series of operations.
    used for query operations
    """
    session = session_cls(**params)
    try:
        for x in initial:
            session.execute(x)
        yield session
    finally:
        session.close()


class Alchemist:
    """

    Provides thread safe named access to sqlalchemy sessions

    first declare your connection::

        ConnectionMgr.connection("default", "sqlite:///simple.db")

    then connect::

        with ConnectionMgr.session("default") as session:
            session.execute(sql.insert(role, {"name": "user"}))
            session.commit()

    """

    connections = {}
    connections_lock = threading.Lock()

    @classmethod
    def connection(cls, using, force=False, initial=None, **kwargs):
        """
        creates a cached engine and session_cls for db_url
        kwargs are passed on the create_engine
        """
        from sqlalchemy import create_engine
        from sqlalchemy.orm import sessionmaker
        with cls.connections_lock:
            conn = cls.connections.get(using)
            if conn and not force:
                return conn
            conf, connector = ConnectionManager.get_connector(using)
            url = connector.sql_alchemy_connection_string(conf)
            initial = initial or conf.get('initial') or []
            kwargs.setdefault('pool_pre_ping', True)
            engine = create_engine(url, **kwargs)
            session_cls = sessionmaker(bind=engine)
            conn = engine, session_cls, initial
            cls.connections[using] = conn
        return conn

    @classmethod
    def transactional(cls, using, force=False):
        """
        returns a sql alchemy session for the provided
        database connection that should be defined in waddle
        use the ``force`` flag to force this function to return
        a new connection instead of using the cached connection.
        The semantics of a transactional session are such that at the beginning
        of the context, we will begin a transaction, and once the context
        exits, we will either commit or rollback the transaction.
        """
        conn = cls.connection(using, force=force)
        return transactional(conn[1], conn[2])

    @classmethod
    def read_only(cls, name='default', force=False, **params):
        """
        returns a sql alchemy session for the provided
        database connection that should be defined in waddle
        use the ``force`` flag to force this function to return
        a new connection instead of using the cached connection.
        The semantics of a read only session are that no transactional
        management will occur in the session context.
        """
        conn = cls.connection(name, force=force, **params)
        return read_only(conn[1], conn[2])


def cleanup_cached_connections():
    m = CachingConnectionManager.get_instance()
    m.cleanup()


atexit.register(cleanup_cached_connections)
