"""
Caprail Migration System.

Migrations provide a version-controlled way to manage your database schema.

Usage::

    class CreateUsersTable(Migration):
        async def up(self, schema: Schema) -> None:
            await schema.create("users", lambda table: [
                table.id(),
                table.string("name"),
                table.string("email").unique(),
                table.string("password"),
                table.timestamps(),
            ])

        async def down(self, schema: Schema) -> None:
            await schema.drop_if_exists("users")
"""
from __future__ import annotations

from abc import ABC, abstractmethod
from datetime import datetime
from pathlib import Path


class ColumnDefinition:
    """Fluent column builder."""

    def __init__(self, col_type: str, name: str, **kwargs) -> None:
        self.col_type = col_type
        self.name = name
        self._nullable = False
        self._default = None
        self._unique = False
        self._index = False
        self._unsigned = False
        self._length = kwargs.get("length")
        self._precision = kwargs.get("precision")
        self._scale = kwargs.get("scale")

    def nullable(self) -> "ColumnDefinition":
        self._nullable = True
        return self

    def default(self, value) -> "ColumnDefinition":
        self._default = value
        return self

    def unique(self) -> "ColumnDefinition":
        self._unique = True
        return self

    def index(self) -> "ColumnDefinition":
        self._index = True
        return self

    def unsigned(self) -> "ColumnDefinition":
        self._unsigned = True
        return self

    def to_sqlalchemy(self):
        import sqlalchemy as sa
        type_map = {
            "id": sa.Integer,
            "big_integer": sa.BigInteger,
            "integer": sa.Integer,
            "small_integer": sa.SmallInteger,
            "string": sa.String,
            "text": sa.Text,
            "long_text": sa.Text,
            "boolean": sa.Boolean,
            "float": sa.Float,
            "decimal": sa.Numeric,
            "date": sa.Date,
            "datetime": sa.DateTime,
            "timestamp": sa.DateTime,
            "time": sa.Time,
            "json": sa.JSON,
            "binary": sa.LargeBinary,
            "uuid": sa.String(36),
        }
        sa_type = type_map.get(self.col_type, sa.String)
        if callable(sa_type) and self._length and self.col_type == "string":
            sa_type = sa.String(self._length or 255)
        elif callable(sa_type):
            sa_type = sa_type()

        kwargs = {
            "nullable": self._nullable,
        }
        if self._default is not None:
            kwargs["server_default"] = str(self._default)
        col = sa.Column(self.name, sa_type, **kwargs)
        return col


class Blueprint:
    """
    Fluent table definition builder passed to the schema callback.
    """

    def __init__(self) -> None:
        self._columns: list[ColumnDefinition] = []
        self._primary_keys: list[str] = []

    def id(self, name: str = "id") -> ColumnDefinition:
        col = ColumnDefinition("id", name)
        import sqlalchemy as sa
        self._columns.append(col)
        self._primary_keys.append(name)
        return col

    def big_increments(self, name: str = "id") -> ColumnDefinition:
        col = ColumnDefinition("big_integer", name)
        self._columns.append(col)
        self._primary_keys.append(name)
        return col

    def string(self, name: str, length: int = 255) -> ColumnDefinition:
        col = ColumnDefinition("string", name, length=length)
        self._columns.append(col)
        return col

    def text(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("text", name)
        self._columns.append(col)
        return col

    def long_text(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("long_text", name)
        self._columns.append(col)
        return col

    def integer(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("integer", name)
        self._columns.append(col)
        return col

    def big_integer(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("big_integer", name)
        self._columns.append(col)
        return col

    def boolean(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("boolean", name)
        self._columns.append(col)
        return col

    def float(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("float", name)
        self._columns.append(col)
        return col

    def decimal(self, name: str, precision: int = 8, scale: int = 2) -> ColumnDefinition:
        col = ColumnDefinition("decimal", name, precision=precision, scale=scale)
        self._columns.append(col)
        return col

    def date(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("date", name)
        self._columns.append(col)
        return col

    def datetime(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("datetime", name)
        self._columns.append(col)
        return col

    def timestamp(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("timestamp", name)
        self._columns.append(col)
        return col

    def timestamps(self) -> None:
        self._columns.append(ColumnDefinition("timestamp", "created_at").nullable())
        self._columns.append(ColumnDefinition("timestamp", "updated_at").nullable())

    def soft_deletes(self) -> None:
        self._columns.append(ColumnDefinition("timestamp", "deleted_at").nullable())

    def json(self, name: str) -> ColumnDefinition:
        col = ColumnDefinition("json", name)
        self._columns.append(col)
        return col

    def uuid(self, name: str = "id") -> ColumnDefinition:
        col = ColumnDefinition("uuid", name)
        self._columns.append(col)
        return col

    def foreign_id(self, name: str) -> ColumnDefinition:
        return self.big_integer(name)

    def foreign(self, column: str):
        return self

    def references(self, column: str):
        return self

    def on(self, table: str):
        return self

    def to_sqlalchemy_table(self, name: str):
        import sqlalchemy as sa
        cols = [c.to_sqlalchemy() for c in self._columns]
        if self._primary_keys:
            pk_cols = [sa.Column(k, sa.Integer, primary_key=True, autoincrement=True)
                       for k in self._primary_keys]
            cols = pk_cols + [c for c in cols if c.name not in self._primary_keys]
        return sa.Table(name, sa.MetaData(), *cols)


class Schema:
    """
    Schema builder — passed to migration up()/down() methods.
    """

    def __init__(self, engine) -> None:
        self._engine = engine

    async def create(self, table_name: str, callback) -> None:
        blueprint = Blueprint()
        callback(blueprint)
        import sqlalchemy as sa
        tbl = blueprint.to_sqlalchemy_table(table_name)
        async with self._engine.begin() as conn:
            await conn.run_sync(tbl.metadata.create_all)

    async def drop(self, table_name: str) -> None:
        import sqlalchemy as sa
        async with self._engine.begin() as conn:
            await conn.run_sync(
                lambda sync_conn: sa.Table(table_name, sa.MetaData()).drop(sync_conn, checkfirst=True)
            )

    async def drop_if_exists(self, table_name: str) -> None:
        await self.drop(table_name)

    async def has_table(self, table_name: str) -> bool:
        import sqlalchemy as sa
        async with self._engine.connect() as conn:
            return await conn.run_sync(
                lambda sync_conn: sa.inspect(sync_conn).has_table(table_name)
            )

    async def rename(self, from_name: str, to_name: str) -> None:
        async with self._engine.begin() as conn:
            await conn.execute(sa.text(f"ALTER TABLE {from_name} RENAME TO {to_name}"))


class Migration(ABC):
    """
    Base class for all database migrations.
    """

    @abstractmethod
    async def up(self, schema: Schema) -> None:
        """Apply the migration."""

    @abstractmethod
    async def down(self, schema: Schema) -> None:
        """Reverse the migration."""


class Migrator:
    """
    Runs migrations, tracks which have been applied, supports rollback.
    """

    MIGRATIONS_TABLE = "caprail_migrations"

    def __init__(self, engine) -> None:
        self._engine = engine
        self._schema = Schema(engine)

    async def install(self) -> None:
        """Create the migrations tracking table if it doesn't exist."""
        import sqlalchemy as sa
        meta = sa.MetaData()
        tbl = sa.Table(
            self.MIGRATIONS_TABLE,
            meta,
            sa.Column("id", sa.Integer, primary_key=True, autoincrement=True),
            sa.Column("migration", sa.String(255)),
            sa.Column("batch", sa.Integer),
            sa.Column("run_at", sa.DateTime),
        )
        async with self._engine.begin() as conn:
            await conn.run_sync(meta.create_all)

    async def run(self, migrations_path: Path) -> list[str]:
        """Run all pending migrations. Returns list of migration names run."""
        await self.install()
        ran = await self._get_ran()
        pending = await self._get_pending(migrations_path, ran)
        if not pending:
            return []
        batch = await self._get_next_batch()
        applied = []
        for name, cls in pending:
            instance = cls()
            await instance.up(self._schema)
            await self._record(name, batch)
            applied.append(name)
        return applied

    async def rollback(self, migrations_path: Path, steps: int = 1) -> list[str]:
        """Roll back the last N batches."""
        await self.install()
        last = await self._get_last_batch_migrations(steps)
        rolled_back = []
        classes = await self._load_migration_classes(migrations_path)
        for name in reversed(last):
            cls = classes.get(name)
            if cls:
                instance = cls()
                await instance.down(self._schema)
                await self._delete_record(name)
                rolled_back.append(name)
        return rolled_back

    async def fresh(self, migrations_path: Path) -> list[str]:
        """Drop all tables and re-run all migrations."""
        async with self._engine.begin() as conn:
            await conn.run_sync(lambda c: sa.MetaData().drop_all(c))
        import sqlalchemy as sa
        return await self.run(migrations_path)

    async def _get_ran(self) -> set[str]:
        import sqlalchemy as sa
        async with self._engine.connect() as conn:
            result = await conn.execute(
                sa.text(f"SELECT migration FROM {self.MIGRATIONS_TABLE}")
            )
            return {row[0] for row in result.fetchall()}

    async def _get_pending(self, path: Path, ran: set[str]) -> list[tuple[str, type]]:
        classes = await self._load_migration_classes(path)
        return [(name, cls) for name, cls in sorted(classes.items()) if name not in ran]

    async def _load_migration_classes(self, path: Path) -> dict[str, type]:
        import importlib.util
        classes: dict[str, type] = {}
        if not path.exists():
            return classes
        for f in sorted(path.glob("*.py")):
            if f.name.startswith("_"):
                continue
            spec = importlib.util.spec_from_file_location(f.stem, f)
            if spec and spec.loader:
                module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(module)
                for attr_name in dir(module):
                    attr = getattr(module, attr_name)
                    if (isinstance(attr, type) and issubclass(attr, Migration)
                            and attr is not Migration):
                        classes[f.stem] = attr
        return classes

    async def _record(self, name: str, batch: int) -> None:
        import sqlalchemy as sa
        async with self._engine.begin() as conn:
            await conn.execute(
                sa.text(
                    f"INSERT INTO {self.MIGRATIONS_TABLE} (migration, batch, run_at) "
                    "VALUES (:migration, :batch, :run_at)"
                ),
                {"migration": name, "batch": batch, "run_at": datetime.utcnow()},
            )

    async def _delete_record(self, name: str) -> None:
        import sqlalchemy as sa
        async with self._engine.begin() as conn:
            await conn.execute(
                sa.text(f"DELETE FROM {self.MIGRATIONS_TABLE} WHERE migration = :name"),
                {"name": name},
            )

    async def _get_next_batch(self) -> int:
        import sqlalchemy as sa
        async with self._engine.connect() as conn:
            result = await conn.execute(
                sa.text(f"SELECT MAX(batch) FROM {self.MIGRATIONS_TABLE}")
            )
            val = result.scalar()
            return (val or 0) + 1

    async def _get_last_batch_migrations(self, steps: int) -> list[str]:
        import sqlalchemy as sa
        async with self._engine.connect() as conn:
            result = await conn.execute(
                sa.text(f"SELECT MAX(batch) FROM {self.MIGRATIONS_TABLE}")
            )
            max_batch = result.scalar() or 1
            names = []
            for b in range(max_batch, max(0, max_batch - steps), -1):
                r2 = await conn.execute(
                    sa.text(
                        f"SELECT migration FROM {self.MIGRATIONS_TABLE} WHERE batch = :b ORDER BY id DESC"
                    ),
                    {"b": b},
                )
                names.extend(row[0] for row in r2.fetchall())
            return names
