"""Test pipeline for running alembic migrations and corresponding test hooks in order."""

from __future__ import annotations

import pytest
from alembic import command
from alembic.config import Config
from alembic.script import ScriptDirectory
from sqlalchemy import Engine

from tests.alembic.conftest import ALEMBIC_CONFIG_PATH
from tests.alembic.registry import get_registered_hooks
from tests.alembic.verifier import MigrationVerifier


def _load_revision_sequence() -> tuple[str, ...]:
    """Read the Alembic script directory to produce the linear revision order."""

    script = ScriptDirectory.from_config(Config(str(ALEMBIC_CONFIG_PATH)))
    revisions = list(script.walk_revisions())  # newest -> oldest
    revisions.reverse()
    return tuple(revision.revision for revision in revisions)


REVISION_SEQUENCE: tuple[str, ...] = _load_revision_sequence()
REVISION_PARAMS = [
    pytest.param(revision, id=f"{index:02d}_{revision}")
    for index, revision in enumerate(REVISION_SEQUENCE, start=1)
]


def _test_single_revision(
    revision: str,
    alembic_cfg: Config,
    alembic_engine: Engine,
) -> None:
    """
    Test a single migration revision upgrade.

    For each revision:
    - Migrate to the previous revision
    - Run the before_upgrade hook to seed and validate the state of the DB before the revision
    - Migrate to the current revision
    - Run the after_upgrade hook to validate the state of the DB after the revision
    """
    hooks_map = get_registered_hooks()
    revision_order = list(REVISION_SEQUENCE)

    # Find the previous revision in the chain
    revision_index = revision_order.index(revision)
    previous_revision = (
        revision_order[revision_index - 1] if revision_index > 0 else "base"
    )

    # Migrate up to the previous revision using a shared connection
    with alembic_engine.begin() as conn:
        alembic_cfg.attributes["connection"] = conn
        command.upgrade(alembic_cfg, previous_revision)

    # Run before_upgrade hook if it exists
    hooks = hooks_map.get(revision)
    if hooks and hooks.before_upgrade:
        with alembic_engine.begin() as conn:
            verifier = MigrationVerifier(conn, revision)
            hooks.before_upgrade(verifier)

    # Migrate to the current revision using the same pattern
    with alembic_engine.begin() as conn:
        alembic_cfg.attributes["connection"] = conn
        command.upgrade(alembic_cfg, revision)

    # Run after_upgrade hook if it exists
    if hooks and hooks.after_upgrade:
        with alembic_engine.begin() as conn:
            verifier = MigrationVerifier(conn, revision)
            hooks.after_upgrade(verifier)


@pytest.mark.parametrize("revision", REVISION_PARAMS)
def test_migration_revision(
    revision: str,
    alembic_cfg: Config,
    alembic_engine: Engine,
) -> None:
    _test_single_revision(revision, alembic_cfg, alembic_engine)