from typing import Callable

import sqlalchemy
from alembic import command
from alembic.config import Config
from sqlalchemy.pool import NullPool

from ._AlembicDatabaseInfo import AlembicDatabaseInfo
from ._DatabaseInfo import DatabaseConfig


class PostgresAlembicHandler(object):
    def __init__(
        self,
        database_info: AlembicDatabaseInfo,
    ):
        self.database_info = 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")
            conn.connection.commit()

        self._execute_in_all_databases(inner_func)

    def _execute_in_all_databases(self, func: Callable[[Config], None]):
        alembic_cfg = Config()
        conn_string = _build_conn_string(self.database_info.db_config)
        alembic_cfg.set_main_option("script_location", self.database_info.alembic_path)
        alembic_cfg.set_main_option("version_locations", self.database_info.versions_dir)
        alembic_cfg.set_main_option("path_separator", "os")
        alembic_cfg.set_main_option("sqlalchemy.url", conn_string)

        engine = sqlalchemy.create_engine(
            url=conn_string,
            poolclass=NullPool,
        )
        conn = engine.connect()
        alembic_cfg.attributes["connection"] = conn
        try:
            func(alembic_cfg)
        finally:
            conn.close()
            engine.dispose()


def _build_conn_string(db_conf: DatabaseConfig) -> str:
    driver = _detect_postgres_driver()
    return (
        f"postgresql+{driver}://{db_conf.user}:{db_conf.password}"
        f"@{db_conf.host}:{db_conf.port}/{db_conf.database}"
    )


def _detect_postgres_driver() -> str:
    """Return the SQLAlchemy driver name for the installed Postgres DBAPI.

    Prefers psycopg2 (for backward compatibility with projects that already
    depend on it). Falls back to psycopg v3 if psycopg2 is not available.
    Raises RuntimeError if neither is installed.
    """
    try:
        import psycopg2  # noqa: F401
        return "psycopg2"
    except ImportError:
        pass
    try:
        import psycopg  # noqa: F401
        return "psycopg"
    except ImportError:
        pass
    raise RuntimeError(
        "No PostgreSQL driver found. Install one of the extras:\n"
        "  pip install postgres_alembic_util[psycopg2]   # psycopg2 (v2.x)\n"
        "  pip install postgres_alembic_util[psycopg]    # psycopg (v3.x)"
    )
