import asyncio
import logging
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from contextvars import ContextVar

import pydantic
from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)

_session_ctx: ContextVar[AsyncSession | None] = ContextVar('session', default=None)


class DatabaseConfig(pydantic.BaseModel):
    URL: pydantic.PostgresDsn
    ECHO: bool = False
    ECHO_POOL: bool = False
    POOL_SIZE: int = 50
    MAX_OVERFLOW: int = 10


class DatabaseBase:
    def __init__(self, config: DatabaseConfig, migrations_path: str = 'migration/versions') -> None:
        self.config = config
        self.migrations_path = migrations_path
        self.engine: AsyncEngine | None = None
        self.session_factory: async_sessionmaker[AsyncSession] | None = None
        self._connect_lock = asyncio.Lock()

    async def connect(self) -> None:
        async with self._connect_lock:
            if self.engine is not None:
                logging.warning('Database already connected')
                return

            self.engine = create_async_engine(
                url=str(self.config.URL),
                echo=self.config.ECHO,
                echo_pool=self.config.ECHO_POOL,
                pool_size=self.config.POOL_SIZE,
                max_overflow=self.config.MAX_OVERFLOW,
            )

            self.session_factory = async_sessionmaker(
                bind=self.engine,
                autoflush=False,
                autocommit=False,
                expire_on_commit=False,
            )

            await self.ping()
            await self._check_migrations()

            logging.info('Database initialized')

    async def ping(self) -> None:
        if self.session_factory is None:
            raise RuntimeError('Database not connected. Call connect() first.')

        try:
            async with self.session_factory() as session:
                await session.execute(text('SELECT 1'))
        except Exception as e:
            logging.error(f'Database ping failed: {e}')

    async def close(self) -> None:
        async with self._connect_lock:
            if self.engine is None:
                logging.warning('Database already closed')
                return

            await self.engine.dispose()
            self.engine = None
            self.session_factory = None

            logging.info('Database closed')

    @property
    def session(self) -> AsyncSession:
        session = _session_ctx.get()
        if session is None:
            raise RuntimeError('No active transaction. Use transaction() or begin() first.')
        return session

    async def begin(self) -> None:
        if _session_ctx.get() is not None:
            raise RuntimeError('Transaction already started')

        if self.session_factory is None:
            raise RuntimeError('Database not connected. Call connect() first.')

        session = self.session_factory()
        _session_ctx.set(session)

    async def commit(self) -> None:
        session = _session_ctx.get()
        if session is None:
            raise RuntimeError('No active transaction')

        try:
            await session.commit()
        finally:
            await session.close()
            _session_ctx.set(None)

    async def rollback(self) -> None:
        session = _session_ctx.get()
        if session is None:
            raise RuntimeError('No active transaction')

        try:
            await session.rollback()
        finally:
            await session.close()
            _session_ctx.set(None)

    @asynccontextmanager
    async def transaction(self) -> AsyncGenerator[None]:
        await self.begin()
        try:
            yield
            await self.commit()
        except Exception:
            await self.rollback()
            raise

    async def _check_migrations(self) -> None:
        if self.session_factory is None:
            raise RuntimeError('Database not connected. Call connect() first.')

        async with self.session_factory() as session:
            result = await session.execute(
                text("SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'alembic_version')")
            )
            if not result.scalar():
                logging.error('Migrations not applied. Run: alembic upgrade head')
                return

            result = await session.execute(text('SELECT version_num FROM alembic_version'))
            current_version = result.scalar()

            if not current_version:
                logging.error('No migration version found in database. Run: alembic upgrade head')
                return

            try:
                from alembic.config import Config
                from alembic.script import ScriptDirectory

                alembic_cfg = Config('alembic.ini')
                script = ScriptDirectory.from_config(alembic_cfg)
                head_revision = script.get_current_head()

                if head_revision and current_version != head_revision:
                    logging.warning(
                        f'Database version {current_version} is outdated (latest: {head_revision}). '
                        f'Run: alembic upgrade head'
                    )
            except Exception as e:
                logging.debug(f'Could not check migration head: {e}')