import pytest
import sqlalchemy

from ehrql.backends.emis import EMISBackend
from ehrql.backends.tpp import TPPBackend


def _get_select_all_query(request, backend):
    try:
        ql_table = request.function._table
    except AttributeError:  # pragma: no cover
        raise RuntimeError(
            f"Function '{request.function.__name__}' needs the "
            f"`@register_test_for(table)` decorator applied"
        )

    qm_table = ql_table._qm_node
    sql_table = backend.get_table_expression(qm_table.name, qm_table.schema)
    columns = [
        # Using `type_coerce(..., None)` like this strips the type information from the
        # SQLAlchemy column meaning we get back the type that the column actually is in
        # database, not the type we've told SQLAlchemy it is.
        sqlalchemy.type_coerce(column, None).label(column.key)
        for column in sql_table.columns
    ]
    return sqlalchemy.select(*columns)


def _select_all_fn(select_all_query, database):
    def _select_all(*input_data):
        database.setup(*input_data)
        with database.engine().connect() as connection:
            results = connection.execute(select_all_query)
            return sorted(
                [row._asdict() for row in results], key=lambda x: x["patient_id"]
            )

    return _select_all


@pytest.fixture
def select_all_emis(request, trino_database):
    select_all_query = _get_select_all_query(request, EMISBackend())
    return _select_all_fn(select_all_query, trino_database)


@pytest.fixture
def select_all_tpp(request, mssql_database):
    backend = TPPBackend(config={"TEMP_DATABASE_NAME": "temp_tables"})
    select_all_query = _get_select_all_query(request, backend)
    return _select_all_fn(select_all_query, mssql_database)