from __future__ import annotations

import contextlib
import logging
import typing as t

from pogo_core import error
from pogo_core.migration import Migration, topological_sort
from pogo_core.util import sql

if t.TYPE_CHECKING:
    from pathlib import Path

    import asyncpg

    from pogo_core.types import Logger

logger_ = logging.getLogger(__name__)


@contextlib.asynccontextmanager
async def transaction(db: asyncpg.Connection, migration: Migration) -> t.AsyncIterator[None]:
    tr = None
    if migration.use_transaction:
        tr = db.transaction()
        await tr.start()

    try:
        yield
    except Exception:
        if tr:  # pragma: no cover
            await tr.rollback()
        raise
    else:
        if tr:
            await tr.commit()


async def apply(
    db: asyncpg.Connection,
    migrations_dir: Path,
    *,
    schema_name: str,
    logger: Logger | logging.Logger | None = None,
) -> None:
    logger = logger or logger_
    await sql.ensure_pogo_sync(db)
    migrations = await sql.read_migrations(migrations_dir, db, schema_name=schema_name)
    migrations = topological_sort([m.load() for m in migrations])

    for migration in migrations:
        try:
            migration.load()
            if not migration.applied:
                logger.warning("Applying %s", migration.id)
                async with transaction(db, migration):
                    await migration.apply(db)
                    await sql.migration_applied(db, migration.id, migration.hash, schema_name=schema_name)
        except Exception as e:  # noqa: PERF203
            msg = f"Failed to apply {migration.id}"
            raise error.BadMigrationError(msg) from e


async def rollback(
    db: asyncpg.Connection,
    migrations_dir: Path,
    *,
    schema_name: str,
    count: int | None = None,
    logger: Logger | logging.Logger | None = None,
) -> None:
    logger = logger or logger_
    await sql.ensure_pogo_sync(db)
    migrations = await sql.read_migrations(migrations_dir, db, schema_name=schema_name)
    migrations = reversed(list(topological_sort([m.load() for m in migrations])))

    i = 0
    for migration in migrations:
        try:
            migration.load()
            if migration.applied and (count is None or i < count):
                logger.warning("Rolling back %s", migration.id)

                async with transaction(db, migration):
                    await migration.rollback(db)
                    await sql.migration_unapplied(db, migration.id, schema_name=schema_name)
                    i += 1
        except Exception as e:  # noqa: PERF203
            msg = f"Failed to rollback {migration.id}"
            raise error.BadMigrationError(msg) from e
