"""Alembic auto-migrate on startup service."""

import logging
import time

from alembic import command as alembic_command
from alembic.autogenerate import api as ag_api
from alembic.config import Config
from alembic.runtime.migration import MigrationContext
from pydantic import BaseModel, DirectoryPath, Field, FilePath
from sqlalchemy.engine import make_url
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlmodel import SQLModel

logger = logging.getLogger(__name__)


def _sync_url_from_async_engine(db_engine: AsyncEngine) -> str:
    """Map an async driver URL to a sync driver URL for Alembic."""
    url = make_url(str(db_engine.url))
    drv = url.drivername

    # Common mappings
    if drv.endswith("+aiosqlite"):
        url = url.set(drivername="sqlite+pysqlite")  # or just "sqlite"
    elif drv.endswith("+asyncpg"):
        # Use psycopg (SQLAlchemy 2.x) or psycopg2 if that's your stack
        url = url.set(drivername="postgresql+psycopg")
    elif drv.endswith("+asyncmy"):
        url = url.set(drivername="mysql+pymysql")
    elif drv.endswith("+aioodbc"):
        # Pick an ODBC driver that fits your environment if you ever use it
        url = url.set(drivername=url.drivername.split("+", 1)[0])  # e.g., "mssql"
    else:
        # If it's already sync (or unknown), leave it as is
        pass

    return str(url)


class AlembicAutoMigrate(BaseModel):
    """Service to run alembic migrations on app startup, with optional autogeneration."""

    # Locations
    ini_path: FilePath = Field(..., description="Path to alembic.ini (auto-resolve if None).")
    script_location: DirectoryPath = Field(..., description="Alembic 'script_location' folder (auto if None).")

    # Behavior
    autogenerate: bool = Field(default=False, description="Whether we may create a new revision automatically.")
    migration_message: str = Field(default="autogenerated at startup", description="Message for new revisions.")
    compare_type: bool = Field(default=True)
    compare_server_default: bool = Field(default=True)
    include_schemas: bool = Field(default=False)

    # Concurrency safety (Postgres only)
    use_advisory_lock: bool = Field(default=True, description="Use pg_advisory_lock to serialize migrations.")
    advisory_lock_key: int = Field(default=792_002, description="Arbitrary lock key; choose a unique one.")

    # ---------- Public API ----------
    async def run(self, db_engine: AsyncEngine) -> str | None:
        """Run migrations, possibly creating a new revision."""
        sync_url = _sync_url_from_async_engine(db_engine)
        cfg = self._build_cfg(sync_url=sync_url)
        created: str | None = None

        # preamble
        logger.info("starting (env=%s, policy=%s, url=%s, script_location=%s)")

        t0 = time.perf_counter()

        # Pre revision head upgrade (ensures new revisions CAN be created)
        # otherwise you will get alembic.util.exc.CommandError: Target database is not up to date.
        logger.info("upgrading to head...")
        alembic_command.upgrade(cfg, "head")
        logger.info("upgrade complete")

        # explain autogenerate decision
        has_new_revisions = False
        if self.autogenerate:
            logger.debug("checking for metadata diffs...")

            has_diffs = await self._metadata_has_diffs(db_engine)
            logger.info("metadata diffs detected=%s", has_diffs)

            if has_diffs:
                # Alembic command APIs are sync; run in a worker thread
                created = alembic_command.revision(cfg, message=self.migration_message, autogenerate=True)
                if created:
                    has_new_revisions = True
                logger.info("created new revision id=%s", created)
            else:
                logger.info("no new revision created")
        else:
            logger.info("autogenerate skipped (disabled via settings)")

        # Final upgrade to head (if we created a new revision)
        if has_new_revisions:
            logger.info("upgrading to head...")
            alembic_command.upgrade(cfg, "head")
            logger.info("upgrade complete")

        t1 = time.perf_counter()
        logger.info("done in %.2fs (created=%s)", t1 - t0, bool(created))
        return created

    # ---------- Internals ----------
    def _build_cfg(self, *, sync_url: str) -> Config:
        cfg = Config(self.ini_path)
        cfg.set_main_option("script_location", str(self.script_location))
        cfg.set_main_option("sqlalchemy.url", str(sync_url))  # <- SYNC URL here
        return cfg

    async def _metadata_has_diffs(self, aengine: AsyncEngine) -> list:
        """Check if the current metadata has diffs compared to the database schema."""
        async with aengine.begin() as aconn:

            def _compare(sync_conn):
                mc = MigrationContext.configure(
                    sync_conn,
                    opts={
                        "compare_type": self.compare_type,
                        "compare_server_default": self.compare_server_default,
                        "include_schemas": self.include_schemas,
                    },
                )
                diffs = ag_api.compare_metadata(mc, SQLModel.metadata)
                return diffs

            return await aconn.run_sync(_compare)

    def _create_revision(self, cfg: Config):
        """Create a new revision, returning its ID or None."""
        return
