import asyncio
import os
from logging.config import fileConfig
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse

from sqlalchemy import pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import create_async_engine

from alembic import context

_ASYNCPG_UNSUPPORTED = {"sslmode", "channel_binding"}


def _build_async_url(raw: str) -> str:
    url = raw.replace("postgresql://", "postgresql+asyncpg://", 1)
    parsed = urlparse(url)
    qs = parse_qs(parsed.query, keep_blank_values=True)
    for param in _ASYNCPG_UNSUPPORTED:
        qs.pop(param, None)
    cleaned = parsed._replace(query=urlencode(qs, doseq=True))
    return urlunparse(cleaned)


# Import all models so SQLModel.metadata is populated
from sqlmodel import SQLModel

from app.models import *  # noqa: F401, F403

config = context.config

# Alembic reads sqlalchemy.url from alembic.ini, but we override with env var here
# to avoid embedding secrets in the ini file
direct_url = os.environ.get("NEON_DATABASE_URL_DIRECT")
if direct_url:
    config.set_main_option("sqlalchemy.url", direct_url)

if config.config_file_name is not None:
    fileConfig(config.config_file_name)

target_metadata = SQLModel.metadata


def run_migrations_offline() -> None:
    url = config.get_main_option("sqlalchemy.url")
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )
    with context.begin_transaction():
        context.run_migrations()


def do_run_migrations(connection: Connection) -> None:
    # Pre-create alembic_version with a 64-char version_num column so revision
    # IDs like `0012_add_conversation_session_metadata` fit. Alembic's default
    # is VARCHAR(32) and its `version_num_length` env kwarg is unreliable —
    # it's evaluated AFTER alembic_version has already been auto-created on a
    # fresh DB. CREATE TABLE IF NOT EXISTS is idempotent and a no-op when the
    # table already exists.
    connection.exec_driver_sql(
        "CREATE TABLE IF NOT EXISTS alembic_version ("
        "  version_num VARCHAR(64) NOT NULL,"
        "  CONSTRAINT alembic_version_pkc PRIMARY KEY (version_num)"
        ")"
    )
    context.configure(
        connection=connection,
        target_metadata=target_metadata,
        version_num_length=64,
    )
    with context.begin_transaction():
        context.run_migrations()


async def run_async_migrations() -> None:
    raw_url = (
        os.environ.get("NEON_DATABASE_URL_DIRECT") or config.get_main_option("sqlalchemy.url") or ""
    )
    url = _build_async_url(raw_url)
    connectable = create_async_engine(url, poolclass=pool.NullPool)
    async with connectable.connect() as connection:
        await connection.run_sync(do_run_migrations)
        await connection.commit()
    await connectable.dispose()


def run_migrations_online() -> None:
    asyncio.run(run_async_migrations())


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()