import pydantic
from tortoise import Tortoise


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) -> None:
        self.config = config

    async def connect(self) -> None:
        await Tortoise.init(
            db_url=str(self.config.URL),
            modules={'models': ['src.domain']},
            use_tz=True,
            timezone='UTC',
        )

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

    async def close(self) -> None:
        await Tortoise.close_connections()

    async def ping(self) -> None:
        from tortoise import connections

        conn = connections.get('default')
        try:
            await conn.execute_query('SELECT 1')
        except Exception as e:
            import logging

            logging.error(f'Database ping failed: {e}')
            raise

    async def _check_migrations(self) -> None:
        import logging

        from tortoise import connections

        conn = connections.get('default')

        # Проверяем существование таблицы alembic_version
        result = await conn.execute_query(
            "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'alembic_version')"
        )
        if not result[1][0]['exists']:
            logging.error('Migrations not applied. Run: alembic upgrade head')
            return

        # Получаем текущую версию
        result = await conn.execute_query('SELECT version_num FROM alembic_version')
        if not result[1]:
            logging.error('No migration version found in database. Run: alembic upgrade head')
            return

        current_version = result[1][0]['version_num']

        # Сравниваем с head версией
        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}')