"""MigrationRunner — execute Migration objects against a live database.

Tracks applied versions in a ``schema_migrations`` table and provides
:meth:`upgrade`, :meth:`downgrade`, and :meth:`status` commands.  No Alembic
dependency — uses the DB client's :meth:`execute_raw` directly.

Usage::

    from db_core import create_postgres_client
    from repo_core.migrations.runner import Migration, MigrationRunner
    from repo_core.migrations.registry import get_postgres_migrations

    client = create_postgres_client()
    runner = MigrationRunner(client, get_postgres_migrations())

    async with client:
        await runner.ensure_migrations_table()
        await runner.upgrade()
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from repo_core.migrations.builder import Op


# ---------------------------------------------------------------------------
# Migration
# ---------------------------------------------------------------------------


@dataclass
class Migration:
    """A single versioned schema migration.

    Args:
        version: Unique version string, e.g. ``"tg_account_0003"``.
        up: :class:`~repo_core.migrations.builder.Op` applied on upgrade.
        down: :class:`~repo_core.migrations.builder.Op` applied on rollback.
            Optional — if ``None``, :meth:`rollback` raises ``RuntimeError``.
        description: Human-readable label shown in status output.
        dialect: ``"postgresql"`` or ``"doris"``.
    """

    version: str
    up: Op
    down: Op | None = field(default=None)
    description: str = ""
    dialect: str = "postgresql"

    async def apply(self, client: Any) -> None:
        """Execute the :attr:`up` operation against *client*."""
        await self.up.execute(client)

    async def rollback(self, client: Any) -> None:
        """Execute the :attr:`down` operation against *client*.

        Raises:
            RuntimeError: If no :attr:`down` operation is defined.
        """
        if self.down is None:
            raise RuntimeError(
                f"Migration '{self.version}' has no rollback (down) operation defined."
            )
        await self.down.execute(client)


# ---------------------------------------------------------------------------
# MigrationRunner
# ---------------------------------------------------------------------------


class MigrationRunner:
    """Applies and tracks migrations against a database client.

    Args:
        client: A connected :class:`~db_core.clients.base.DBClient` instance.
        migrations: Ordered list of :class:`Migration` objects to manage.
            Pass the full list (both applied and pending) — the runner checks
            the ``schema_migrations`` table to determine what has already run.
        dialect: ``"postgresql"`` or ``"doris"``.  Controls the DDL used for
            the ``schema_migrations`` tracking table and identifier quoting.
    """

    MIGRATIONS_TABLE = "schema_migrations"

    def __init__(
        self,
        client: Any,
        migrations: list[Migration],
        *,
        dialect: str = "postgresql",
    ) -> None:
        self._client = client
        self._migrations = migrations
        self._dialect = dialect

    # ------------------------------------------------------------------
    # Bootstrap
    # ------------------------------------------------------------------

    async def ensure_migrations_table(self) -> None:
        """Create the ``schema_migrations`` tracking table if it does not exist."""
        if self._dialect == "doris":
            sql = (
                f"CREATE TABLE IF NOT EXISTS `{self.MIGRATIONS_TABLE}` (\n"
                f"  `version` VARCHAR(255),\n"
                f"  `applied_at` DATETIME DEFAULT CURRENT_TIMESTAMP\n"
                f") ENGINE=OLAP\n"
                f"DUPLICATE KEY(`version`)\n"
                f"DISTRIBUTED BY HASH(`version`) BUCKETS 1\n"
                f'PROPERTIES ("replication_num" = "1")'
            )
        else:
            sql = (
                f"CREATE TABLE IF NOT EXISTS {self.MIGRATIONS_TABLE} (\n"
                f"  version TEXT PRIMARY KEY,\n"
                f"  applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()\n"
                f")"
            )
        await self._client.execute_raw(sql)

    # ------------------------------------------------------------------
    # Internal state helpers
    # ------------------------------------------------------------------

    async def get_applied_versions(self) -> set[str]:
        """Return the set of version strings already recorded in the DB."""
        rows = await self._client.execute_raw(
            f"SELECT version FROM {self.MIGRATIONS_TABLE}"
        )
        if not rows:
            return set()
        return {row["version"] for row in rows}

    async def _record_applied(self, version: str) -> None:
        if self._dialect == "doris":
            await self._client.execute_raw(
                f"INSERT INTO `{self.MIGRATIONS_TABLE}` (`version`) VALUES (:version)",
                {"version": version},
            )
        else:
            await self._client.execute_raw(
                f"INSERT INTO {self.MIGRATIONS_TABLE} (version) VALUES (:version)",
                {"version": version},
            )

    async def _remove_applied(self, version: str) -> None:
        if self._dialect == "doris":
            await self._client.execute_raw(
                f"DELETE FROM `{self.MIGRATIONS_TABLE}` WHERE `version` = :version",
                {"version": version},
            )
        else:
            await self._client.execute_raw(
                f"DELETE FROM {self.MIGRATIONS_TABLE} WHERE version = :version",
                {"version": version},
            )

    # ------------------------------------------------------------------
    # Upgrade / downgrade
    # ------------------------------------------------------------------

    async def upgrade(self) -> None:
        """Apply all pending migrations in registration order."""
        applied = await self.get_applied_versions()
        pending = [m for m in self._migrations if m.version not in applied]

        if not pending:
            print("  All migrations already applied.")
            return

        for migration in pending:
            print(f"  Applying  {migration.version}: {migration.description}")
            await migration.apply(self._client)
            await self._record_applied(migration.version)
            print(f"  Applied   {migration.version}")

    async def downgrade(self, target: str | None = None) -> None:
        """Roll back migrations.

        Args:
            target: If given, roll back all migrations that were applied *after*
                this version (exclusive).  If ``None``, roll back only the last
                applied migration.
        """
        applied = await self.get_applied_versions()
        # Iterate in reverse registration order to undo latest first
        applied_in_order = [m for m in reversed(self._migrations) if m.version in applied]

        if not applied_in_order:
            print("  Nothing to roll back.")
            return

        if target is None:
            to_rollback = [applied_in_order[0]]
        else:
            to_rollback = []
            for m in applied_in_order:
                if m.version == target:
                    break
                to_rollback.append(m)

            if not to_rollback:
                print(
                    f"  Nothing to roll back "
                    f"(target '{target}' may already be the current head)."
                )
                return

        for migration in to_rollback:
            print(f"  Rolling back  {migration.version}: {migration.description}")
            await migration.rollback(self._client)
            await self._remove_applied(migration.version)
            print(f"  Rolled back   {migration.version}")

    # ------------------------------------------------------------------
    # Status
    # ------------------------------------------------------------------

    async def status(self) -> list[dict[str, Any]]:
        """Return the status of every registered migration.

        Returns:
            List of dicts with keys:
            ``version``, ``description``, ``status`` (``"applied"``/``"pending"``),
            ``has_rollback``.
        """
        applied = await self.get_applied_versions()
        return [
            {
                "version": m.version,
                "description": m.description,
                "status": "applied" if m.version in applied else "pending",
                "has_rollback": m.down is not None,
            }
            for m in self._migrations
        ]
