from abc import abstractmethod
from contextlib import asynccontextmanager

from iceaxe.migrations.migrator import Migrator
from iceaxe.session import DBConnection


class MigrationRevisionBase:
    """
    Base class for all revisions. This class is most often automatically
    generated by the `migrate` CLI command, which automatically determines
    the proper up and down revisions for your migration class.

    Once added to your project, you can modify the up/down migration methods
    however you see fit.

    """

    # up and down revision are both set, except for the initial revision
    # where down_revision is None
    up_revision: str
    down_revision: str | None

    use_transaction: bool = True
    """
    Disables the transaction for the current migration. Only do this if you're
    confident that the migration will succeed on the first try, or is otherwise
    independent so it can be run multiple times.

    This can speed up migrations, and in some cases might be even fully required for your
    production database to avoid deadlocks when interacting with hot tables.

    """

    async def _handle_up(self, db_connection: DBConnection):
        """
        Internal method to handle the up migration.
        """
        # Isolated migrator context just for this migration
        async with self._optional_transaction(db_connection):
            migrator = Migrator(db_connection)
            await self.up(migrator)
            await migrator.set_active_revision(self.up_revision)

    async def _handle_down(self, db_connection: DBConnection):
        """
        Internal method to handle the down migration.
        """
        async with self._optional_transaction(db_connection):
            migrator = Migrator(db_connection)
            await self.down(migrator)
            await migrator.set_active_revision(self.down_revision)

    @asynccontextmanager
    async def _optional_transaction(self, db_connection: DBConnection):
        if self.use_transaction:
            async with db_connection.transaction():
                yield
        else:
            yield

    @abstractmethod
    async def up(self, migrator: Migrator):
        """
        Perform the migration "up" action. This converts your old database schema to the new
        schema that's found in your code definition. Add any other migration rules to your
        data in this function as well.

        It's good practice to make sure any up migrations don't immediately drop data. Instead,
        consider moving to a temporary table.

        Support both raw SQL execution and helper functions to manipulate the tables and
        columns that you have defined. See the Migrator class for more information.

        """
        pass

    @abstractmethod
    async def down(self, migrator: Migrator):
        """
        Perform the migration "down" action. This converts the current database state to the
        previous snapshot. It should be used in any automatic rollback pipelines if you
        had a feature that was rolled out and then rolled back.

        """
        pass
