"""
cortex/storage/migrations/runner.py — Migration runner for Cortex SQLite schema.

Discovers migration modules in this package, compares VERSION against
the schema_version stored in cortex_config, and applies any pending
migrations in order.

Each migration is applied inside its own transaction; if a migration
fails the transaction is rolled back and the error is re-raised, leaving
the database at the last successfully-applied version.

Idempotency: migrations whose VERSION <= current schema_version are
skipped, so calling apply_all() multiple times is safe.
"""
from __future__ import annotations

import importlib
import logging
import pkgutil
from pathlib import Path
from typing import Any, List, Optional

import aiosqlite

logger = logging.getLogger(__name__)

_MIGRATIONS_PKG = "cortex.storage.migrations"
_MIGRATIONS_DIR = Path(__file__).parent


class MigrationRunner:
    """Apply pending schema migrations to a live aiosqlite connection.

    Parameters
    ----------
    conn:
        An open :class:`aiosqlite.Connection`.  The runner does NOT open
        or close the connection; that remains the caller's responsibility.
    """

    def __init__(self, conn: aiosqlite.Connection) -> None:
        self._conn = conn

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def get_current_version(self) -> int:
        """Read schema_version from cortex_config, return 0 if absent."""
        try:
            async with self._conn.execute(
                "SELECT value FROM cortex_config WHERE key = 'schema_version'"
            ) as cur:
                row = await cur.fetchone()
                return int(row[0]) if row else 0
        except Exception as exc:
            # "no such table" is expected on first run (before v3_initial)
            if "no such table" in str(exc):
                return 0
            raise  # Re-raise real errors (corruption, locks, etc.)

    def get_pending_migrations(self) -> List[Any]:
        """Return all migration modules with VERSION > current_version, sorted ascending.

        This is a synchronous method — it only scans the package on disk,
        it does not touch the database.  Call :meth:`get_current_version`
        first to know the cutoff.

        Returns a list of imported module objects, each guaranteed to have
        ``VERSION``, ``DESCRIPTION``, and ``SQL_UP`` attributes.
        """
        return self._load_all_migrations()

    async def get_pending_since(self, current_version: int) -> List[Any]:
        """Return migration modules with VERSION > *current_version*, sorted ascending."""
        all_migrations = self._load_all_migrations()
        return [m for m in all_migrations if m.VERSION > current_version]

    async def apply_all(self) -> int:
        """Apply all pending migrations in version order.

        Returns
        -------
        int
            Number of migrations applied.  0 means schema was already
            up-to-date.
        """
        current = await self.get_current_version()
        pending = await self.get_pending_since(current)

        if not pending:
            logger.debug("Schema at v%d — no migrations pending.", current)
            return 0

        applied = 0
        for migration in pending:
            success = await self.apply_one(migration)
            if success:
                applied += 1
                logger.info(
                    "Migration v%d applied: %s",
                    migration.VERSION,
                    migration.DESCRIPTION,
                )

        return applied

    async def apply_one(self, migration: Any) -> bool:
        """Apply a single migration inside an explicit transaction.

        On failure the transaction is rolled back and the exception is
        re-raised.  The caller is responsible for deciding whether to
        abort the whole sequence or skip to the next migration.

        Parameters
        ----------
        migration:
            An imported migration module with VERSION, DESCRIPTION,
            SQL_UP (and optionally SQL_DOWN) attributes.

        Returns
        -------
        bool
            True on success.
        """
        version = migration.VERSION
        description = migration.DESCRIPTION
        sql_up = migration.SQL_UP.strip()

        logger.debug("Applying migration v%d: %s", version, description)

        # Strategy:
        # - Run SQL_UP via executescript so that DDL statements using
        #   DEFAULT (datetime('now')) work correctly on all Python versions.
        #   executescript auto-commits any open transaction before running.
        # - Update schema_version in a separate explicit transaction that
        #   we can roll back if the version-bookkeeping step itself fails.
        #
        # Note: SQLite DDL (CREATE TABLE, ALTER TABLE) is effectively
        # auto-committed at the C level; ROLLBACK after failed DDL is
        # best-effort.  For DML-only migrations the rollback is reliable.
        try:
            # executescript requires a plain str; it handles multi-statement
            # DDL correctly including dynamic DEFAULT expressions.
            await self._conn.executescript(sql_up)
        except Exception as exc:
            logger.error(
                "Migration v%d SQL_UP FAILED (%s): %s", version, description, exc
            )
            raise

        # Persist the new schema_version in an explicit transaction so the
        # bookkeeping step is atomic and rollback-safe.
        try:
            await self._conn.execute("BEGIN")
            try:
                await self._conn.execute(
                    """
                    INSERT INTO cortex_config (key, value, updated_at)
                    VALUES ('schema_version', ?, datetime('now'))
                    ON CONFLICT(key) DO UPDATE
                        SET value = excluded.value,
                            updated_at = excluded.updated_at
                    """,
                    (str(version),),
                )
                await self._conn.execute("COMMIT")
            except Exception:
                await self._conn.execute("ROLLBACK")
                raise
        except Exception as exc:
            logger.error(
                "Migration v%d schema_version update FAILED (%s): %s",
                version,
                description,
                exc,
            )
            raise

        return True

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

    @staticmethod
    def _load_all_migrations() -> List[Any]:
        """Import all *migration modules* in this package and return them sorted by VERSION.

        A migration module is any sub-module inside ``cortex.storage.migrations``
        whose name starts with ``v`` and that exposes a ``VERSION`` integer
        attribute.  The ``__init__`` and ``runner`` modules are skipped.
        """
        migrations = []
        skipped_names = {"__init__", "runner"}

        for module_info in pkgutil.iter_modules([str(_MIGRATIONS_DIR)]):
            name = module_info.name
            if name in skipped_names:
                continue
            if not name.startswith("v"):
                continue

            full_name = f"{_MIGRATIONS_PKG}.{name}"
            try:
                mod = importlib.import_module(full_name)
            except ImportError as exc:
                logger.warning("Could not import migration module %s: %s", full_name, exc)
                continue

            # Basic validation
            for attr in ("VERSION", "DESCRIPTION", "SQL_UP"):
                if not hasattr(mod, attr):
                    logger.warning(
                        "Migration module %s is missing attribute '%s' — skipping.",
                        full_name,
                        attr,
                    )
                    break
            else:
                migrations.append(mod)

        migrations.sort(key=lambda m: m.VERSION)
        return migrations
