import os

import alembic
import alembic.command
import alembic.config
import pytest

from smb.common.pgswim.pytest.plugin import (  # noqa
    _database_url_env_var,
    _pgswim_engine_cls,
    _pgswim_stripped_database_url,
    _pgswim_transaction_engine,
    con,
    db,
    pgswim_database_url,
    pgswim_wait_for_db,
    pytest_addoption as base_pytest_addoption,
    pytest_configure as base_pytest_configure,
)

_alembic_ini_path_env_var = "ALEMBIC_INI_PATH"


def pytest_addoption(parser):
    base_pytest_addoption(parser)
    parser.addoption(
        "--alembic-ini-path",
        help=(
            "path to alembic.ini file "
            f"(could be set via {_alembic_ini_path_env_var} environment variable too)"
        ),
        default="alembic.ini",
    )


def pytest_configure(config):
    base_pytest_configure(config)
    config.option.alembic_ini_path = os.getenv(
        _alembic_ini_path_env_var, config.option.alembic_ini_path
    )


@pytest.fixture(scope="session")
def pgswim_alembic_ini_path(pytestconfig):
    return pytestconfig.getoption("--alembic-ini-path")


@pytest.fixture(scope="session")
def pgswim_migrate_db(
    pgswim_alembic_ini_path, _pgswim_stripped_database_url, pgswim_wait_for_db  # noqa
):
    cfg = alembic.config.Config(pgswim_alembic_ini_path)
    cfg.set_main_option("sqlalchemy.url", _pgswim_stripped_database_url)

    alembic.command.upgrade(cfg, "head")
    yield
    alembic.command.downgrade(cfg, "base")


@pytest.fixture
async def _pgswim_real_engine(
    _pgswim_engine_cls,  # noqa
    pgswim_alembic_ini_path,  # noqa
    pgswim_database_url,  # noqa
    _pgswim_stripped_database_url,  # noqa
):
    cfg = alembic.config.Config(pgswim_alembic_ini_path)
    cfg.set_main_option("sqlalchemy.url", _pgswim_stripped_database_url)
    _db = await _pgswim_engine_cls.create(pgswim_database_url, recalc_interval=0)

    yield _db

    await _db.close()

    alembic.command.downgrade(cfg, "base")
    alembic.command.upgrade(cfg, "head")
