import sqlite3
from typing import (
    Union,
    Dict,
    Callable,
    Sequence,
)

import sqlalchemy
from alembic import command
from alembic.config import Config
from sqlalchemy import NullPool
from ._AlembicDatabaseInfo import AlembicDatabaseInfo


class SqliteAlembicHandler(object):
    def __init__(
        self,
        database_infos: Union[AlembicDatabaseInfo, Sequence[AlembicDatabaseInfo]],
    ):
        self.database_infos: Dict[str, AlembicDatabaseInfo] = {}
        if not hasattr(database_infos, '__iter__'):
            database_infos = [database_infos]
        for database_info in database_infos:
            self.database_infos[database_info.name] = database_info

    def upgrade_to_head(self):
        def inner_func(alembic_cfg: Config):
            command.upgrade(alembic_cfg, "head")

        self._execute_in_all_databases(inner_func)

    def downgrade_to_base(self):
        def inner_func(alembic_cfg: Config):
            command.downgrade(alembic_cfg, "base")
            conn = alembic_cfg.attributes["connection"]
            conn.connection.cursor().execute("drop table alembic_version")

        self._execute_in_all_databases(inner_func)

    def _execute_in_all_databases(self, func: Callable[[Config], None]):
        for database_info in self.database_infos.values():
            for conn_string in database_info.conn_strings:
                alembic_cfg = Config()
                alembic_cfg.set_main_option("script_location", database_info.alembic_path)
                alembic_cfg.set_main_option("version_locations", database_info.versions_dir)
                alembic_cfg.set_main_option("sqlalchemy.url", conn_string)

                engine = sqlalchemy.create_engine(
                    url="sqlite:///:memory:",
                    creator=lambda cs=conn_string: sqlite3.connect(database=cs, uri=True),
                    poolclass=NullPool,
                )
                conn = engine.connect()
                alembic_cfg.attributes["connection"] = conn
                try:
                    func(alembic_cfg)
                finally:
                    conn.close()
