"""Database connection and session management."""

from __future__ import annotations

import asyncio
from pathlib import Path
from typing import TYPE_CHECKING

from sqlalchemy import text
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
)
from sqlalchemy.ext.asyncio import (
    create_async_engine as sa_create_async_engine,
)

from noteflow.infrastructure.logging import get_logger, log_timing

if TYPE_CHECKING:
    from noteflow.config import Settings

logger = get_logger(__name__)


def _mask_database_url(url: str) -> str:
    """Mask password in database URL for safe logging.

    Args:
        url: Database URL potentially containing password.

    Returns:
        URL with password replaced by '***'.
    """
    # Handle postgresql+asyncpg://user:password@host:port/db format
    if "://" in url and "@" in url:
        # Split at :// to get scheme and rest
        scheme, rest = url.split("://", 1)
        if "@" in rest:
            # user:password@host:port/db
            auth_part, host_part = rest.rsplit("@", 1)
            if ":" in auth_part:
                user, _ = auth_part.split(":", 1)
                return f"{scheme}://{user}:***@{host_part}"
    return url


def create_async_engine(settings: Settings) -> AsyncEngine:
    """Create an async SQLAlchemy engine.

    Args:
        settings: Application settings with database URL.

    Returns:
        Configured async engine.
    """
    # Mask password in URL for logging
    masked_url = _mask_database_url(settings.database_url_str)
    with log_timing(
        "database_engine_create",
        pool_size=settings.db_pool_size,
        echo=settings.db_echo,
    ):
        engine = sa_create_async_engine(
            settings.database_url_str,
            pool_size=settings.db_pool_size,
            echo=settings.db_echo,
            pool_pre_ping=True,  # Verify connections before use
        )
    logger.info(
        "database_engine_created",
        url=masked_url,
        pool_size=settings.db_pool_size,
    )
    return engine


def get_async_session_factory(
    engine: AsyncEngine,
) -> async_sessionmaker[AsyncSession]:
    """Create an async session factory.

    Args:
        engine: SQLAlchemy async engine.

    Returns:
        Session factory for creating async sessions.
    """
    return async_sessionmaker(
        engine,
        class_=AsyncSession,
        expire_on_commit=False,
        autocommit=False,
        autoflush=False,
    )


def create_async_session_factory(
    database_url: str,
    pool_size: int = 5,
    echo: bool = False,
) -> async_sessionmaker[AsyncSession]:
    """Create an async session factory from a database URL string.

    Convenience function for creating a session factory directly from a URL.

    Args:
        database_url: PostgreSQL database URL.
        pool_size: Connection pool size.
        echo: Enable SQL echo logging.

    Returns:
        Async session factory.
    """
    _engine, factory = create_engine_and_session_factory(database_url, pool_size, echo)
    return factory


def create_engine_and_session_factory(
    database_url: str,
    pool_size: int = 5,
    echo: bool = False,
) -> tuple[AsyncEngine, async_sessionmaker[AsyncSession]]:
    """Create an async engine and session factory from a database URL string.

    Use this when you need access to the engine for disposal during shutdown.

    Args:
        database_url: PostgreSQL database URL.
        pool_size: Connection pool size.
        echo: Enable SQL echo logging.

    Returns:
        Tuple of (engine, session_factory) for lifecycle management.
    """
    masked_url = _mask_database_url(database_url)
    with log_timing(
        "database_engine_and_factory_create",
        pool_size=pool_size,
        echo=echo,
    ):
        engine = sa_create_async_engine(
            database_url,
            pool_size=pool_size,
            echo=echo,
            pool_pre_ping=True,
        )
        factory = async_sessionmaker(
            engine,
            class_=AsyncSession,
            expire_on_commit=False,
            autocommit=False,
            autoflush=False,
        )
    logger.info(
        "database_engine_and_factory_created",
        url=masked_url,
        pool_size=pool_size,
    )
    return engine, factory


async def check_schema_exists(session_factory: async_sessionmaker[AsyncSession]) -> bool:
    """Check if the noteflow schema and core tables exist.

    Verifies that critical tables (meetings, segments, diarization_jobs) exist.
    This ensures the schema is complete, not just partially initialized.

    Args:
        session_factory: Async session factory.

    Returns:
        True if the schema and all core tables exist.
    """
    async with session_factory() as session:
        # Check for multiple critical tables to ensure schema is complete
        result = await session.execute(
            text(
                "SELECT COUNT(*) FROM information_schema.tables "
                "WHERE table_schema = 'noteflow' "
                "AND table_name IN ('meetings', 'segments', 'diarization_jobs')"
            )
        )
        count = result.scalar()
        return count == 3 if count is not None else False


def run_migrations(database_url: str) -> None:
    """Run Alembic migrations programmatically.

    Args:
        database_url: Database URL for migrations.
    """
    from alembic import command
    from alembic.config import Config

    # Find alembic.ini relative to project root
    # The database.py is at src/noteflow/infrastructure/persistence/database.py
    # alembic.ini is at the project root
    current_file = Path(__file__)
    project_root = current_file.parent.parent.parent.parent.parent
    alembic_ini = project_root / "alembic.ini"

    if not alembic_ini.exists():
        msg = f"alembic.ini not found at {alembic_ini}"
        raise FileNotFoundError(msg)

    alembic_cfg = Config(str(alembic_ini))

    # Override the database URL
    # Convert to async driver format if needed
    url = database_url
    if url.startswith("postgres://"):
        url = url.replace("postgres://", "postgresql+asyncpg://", 1)
    elif url.startswith("postgresql://") and "+asyncpg" not in url:
        url = url.replace("postgresql://", "postgresql+asyncpg://", 1)

    alembic_cfg.set_main_option("sqlalchemy.url", url)

    logger.info("Running database migrations...")
    command.upgrade(alembic_cfg, "head")
    logger.info("Database migrations complete")


async def check_alembic_version_exists(
    session_factory: async_sessionmaker[AsyncSession],
) -> bool:
    """Check if Alembic version table exists.

    Args:
        session_factory: Async session factory.

    Returns:
        True if the Alembic version table exists.
    """
    async with session_factory() as session:
        result = await session.execute(
            text(
                "SELECT EXISTS ("
                "SELECT 1 FROM information_schema.tables "
                "WHERE table_schema = 'noteflow' AND table_name = 'alembic_version'"
                ")"
            )
        )
        return bool(result.scalar())


def stamp_alembic_version(database_url: str, revision: str = "head") -> None:
    """Stamp Alembic version table without running migrations.

    Used when schema.sql has created tables but Alembic hasn't recorded a version.

    Args:
        database_url: Database URL for migrations.
        revision: Revision to stamp (default: "head").
    """
    from alembic import command
    from alembic.config import Config

    current_file = Path(__file__)
    project_root = current_file.parent.parent.parent.parent.parent
    alembic_ini = project_root / "alembic.ini"

    if not alembic_ini.exists():
        msg = f"alembic.ini not found at {alembic_ini}"
        raise FileNotFoundError(msg)

    alembic_cfg = Config(str(alembic_ini))

    # Override the database URL
    url = database_url
    if url.startswith("postgres://"):
        url = url.replace("postgres://", "postgresql+asyncpg://", 1)
    elif url.startswith("postgresql://") and "+asyncpg" not in url:
        url = url.replace("postgresql://", "postgresql+asyncpg://", 1)

    alembic_cfg.set_main_option("sqlalchemy.url", url)

    logger.info(f"Stamping Alembic version table with revision: {revision}")
    command.stamp(alembic_cfg, revision)
    logger.info("Alembic version table stamped successfully")


async def _count_noteflow_tables(session_factory: async_sessionmaker[AsyncSession]) -> int:
    """Count tables in noteflow schema excluding alembic_version."""
    async with session_factory() as session:
        result = await session.execute(
            text(
                "SELECT COUNT(*) FROM information_schema.tables "
                "WHERE table_schema = 'noteflow' AND table_name != 'alembic_version'"
            )
        )
        return result.scalar() or 0


async def _table_exists(session: AsyncSession, table_name: str) -> bool:
    """Check if a table exists in noteflow schema."""
    result = await session.execute(
        text(
            "SELECT EXISTS ("
            "SELECT 1 FROM information_schema.tables "
            "WHERE table_schema = 'noteflow' AND table_name = :table_name"
            ")"
        ),
        {"table_name": table_name},
    )
    return bool(result.scalar())


async def _create_user_preferences_table(session: AsyncSession) -> bool:
    """Create user_preferences table if missing. Returns True if created."""
    try:
        await session.execute(
            text("""
                CREATE TABLE IF NOT EXISTS noteflow.user_preferences (
                    key VARCHAR(64) PRIMARY KEY,
                    value JSONB NOT NULL DEFAULT '{}',
                    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now()
                )
            """)
        )
        await session.commit()
        return True
    except SQLAlchemyError as e:
        logger.error("Error creating user_preferences table: %s", e, exc_info=True)
        await session.rollback()
        raise


async def _stamp_database_async(database_url: str) -> None:
    """Stamp database with current Alembic head revision."""
    loop = asyncio.get_running_loop()
    await loop.run_in_executor(None, lambda: stamp_alembic_version(database_url))


async def _run_migrations_async(database_url: str) -> None:
    """Run Alembic migrations."""
    loop = asyncio.get_running_loop()
    await loop.run_in_executor(None, lambda: run_migrations(database_url))


async def _handle_tables_without_alembic(
    session_factory: async_sessionmaker[AsyncSession],
    database_url: str,
    table_count: int,
) -> None:
    """Handle case where tables exist but Alembic version doesn't."""
    critical_tables = ["meetings", "segments", "diarization_jobs", "user_preferences"]
    missing_tables: list[str] = []

    async with session_factory() as session:
        for table_name in critical_tables:
            if not await _table_exists(session, table_name):
                missing_tables.append(table_name)

    if missing_tables:
        logger.warning(
            "Tables exist but missing critical tables: %s. Creating missing tables...",
            ", ".join(missing_tables),
        )
        if "user_preferences" in missing_tables:
            async with session_factory() as session:
                if await _create_user_preferences_table(session):
                    logger.info("Successfully created user_preferences table")

        logger.info("Stamping database after creating missing tables...")
        await _stamp_database_async(database_url)
        logger.info("Database schema ready (created missing tables and stamped)")
        return

    # Safety check for user_preferences
    async with session_factory() as session:
        if not await _table_exists(session, "user_preferences"):
            logger.warning("user_preferences table missing despite check, creating it...")
            await _create_user_preferences_table(session)
            logger.info("Created user_preferences table (safety check)")

    logger.info(
        "Tables exist (%d) but Alembic version table missing, stamping database...",
        table_count,
    )
    await _stamp_database_async(database_url)
    logger.info("Database schema ready (stamped from schema.sql)")


async def _handle_alembic_with_tables(
    session_factory: async_sessionmaker[AsyncSession],
    database_url: str,
    table_count: int,
) -> bool:
    """Handle case where Alembic exists and tables exist. Returns True if handled."""
    if table_count == 0:
        return False

    async with session_factory() as session:
        result = await session.execute(
            text("SELECT version_num FROM noteflow.alembic_version LIMIT 1")
        )
        current_revision = result.scalar()

        if current_revision is None:
            logger.warning(
                "Alembic version table exists but is empty. "
                "Tables exist (%d), stamping database to head...",
                table_count,
            )
            await _stamp_database_async(database_url)
            logger.info("Database schema ready (stamped existing tables)")
            return True

    return False


async def ensure_schema_ready(
    session_factory: async_sessionmaker[AsyncSession],
    database_url: str,
) -> None:
    """Ensure database schema exists, running migrations if needed.

    Handles the case where schema.sql has created tables but Alembic hasn't
    recorded a version by stamping the database with the current head revision.

    Args:
        session_factory: Async session factory.
        database_url: Database URL for migrations.
    """
    alembic_version_exists = await check_alembic_version_exists(session_factory)
    table_count = await _count_noteflow_tables(session_factory)

    logger.debug(
        "Schema check: table_count=%d, alembic_version_exists=%s",
        table_count,
        alembic_version_exists,
    )

    # Case 1: Tables exist but no Alembic version (schema.sql created tables)
    if table_count > 0 and not alembic_version_exists:
        await _handle_tables_without_alembic(session_factory, database_url, table_count)
        return

    # Case 2: Alembic version exists
    if alembic_version_exists:
        if await _handle_alembic_with_tables(session_factory, database_url, table_count):
            return

        logger.debug("Alembic version table exists, checking if migrations needed...")
        try:
            await _run_migrations_async(database_url)
            logger.info("Database migrations checked/updated")
        except (SQLAlchemyError, RuntimeError) as e:
            if (
                "already exists" not in str(e).lower()
                and "duplicate" not in str(e).lower()
            ):
                raise
            logger.warning(
                "Migration failed because tables already exist. "
                "Stamping database to head instead...",
            )
            await _stamp_database_async(database_url)
            logger.info("Database schema ready (stamped after migration conflict)")
        return

    # Case 3: Fresh database - no tables and no Alembic version
    logger.info("Fresh database detected, running migrations...")
    await _run_migrations_async(database_url)
    logger.info("Database initialized successfully")
