import os
import time

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

from maps_adv.common.pg_engine import DB

_database_url_env_var = "DATABASE_URL_FOR_TESTS"
_alembic_ini_path_env_var = "ALEMBIC_INI_PATH"


def pytest_addoption(parser):
    parser.addoption(
        "--database-url",
        help=(
            "database url for tests (will use it for rw and ro in tests) "
            f"(could be set via {_database_url_env_var} environment variable too)"
        ),
        required=False,
    )
    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):
    _database_url = os.getenv(_database_url_env_var, config.option.database_url)
    if not _database_url:
        raise ValueError(
            "Don't know how to connect to database! "
            "Set it with --database-url parameter of pytest "
            f"or {_database_url_env_var} environment variable"
        )

    config.option.database_url = _database_url
    config.option.alembic_ini_path = os.getenv(
        _alembic_ini_path_env_var, config.option.alembic_ini_path
    )


@pytest.fixture(scope="session")
def pg_engine_database_url(pytestconfig):
    return pytestconfig.getoption("--database-url")


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


@pytest.fixture
def _pg_engine_db_cls(request):
    if "pg_engine_db_cls" in request.fixturenames:
        return request.getfixturevalue("pg_engine_db_cls")
    return DB


@pytest.fixture(scope="session")
def pg_engine_wait_for_db(pg_engine_database_url):
    waited = 0

    while waited < 100:
        try:
            con = psycopg2.connect(pg_engine_database_url)
        except psycopg2.OperationalError:
            time.sleep(1)
            waited += 1
        else:
            con.close()
            break
    else:
        raise ConnectionRefusedError()


@pytest.fixture(scope="session")
def pg_engine_migrate_db(
    pg_engine_alembic_ini_path, pg_engine_database_url, pg_engine_wait_for_db
):
    cfg = alembic.config.Config(pg_engine_alembic_ini_path)
    cfg.set_main_option("sqlalchemy.url", pg_engine_database_url)

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


@pytest.fixture
def db(request, pg_engine_migrate_db):
    if request.node.get_closest_marker("real_db"):
        return request.getfixturevalue("_pg_engine_real_db")
    return request.getfixturevalue("_pg_engine_transaction_db")


@pytest.fixture
async def _pg_engine_transaction_db(_pg_engine_db_cls, pg_engine_database_url):
    _db = await _pg_engine_db_cls.create(
        pg_engine_database_url, use_single_connection=True
    )

    con = await _db.acquire()
    tr = con.transaction()
    await tr.start()

    yield _db

    await tr.rollback()
    await _db.release(con, force=True)

    await _db.close()


@pytest.fixture
async def _pg_engine_real_db(
    _pg_engine_db_cls, pg_engine_alembic_ini_path, pg_engine_database_url
):
    cfg = alembic.config.Config(pg_engine_alembic_ini_path)
    cfg.set_main_option("sqlalchemy.url", pg_engine_database_url)
    _db = await _pg_engine_db_cls.create(pg_engine_database_url, pg_engine_database_url)

    yield _db

    await _db.close()

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


@pytest.fixture
async def con(db):
    async with db.acquire() as con:
        yield con
