import asyncio
from contextlib import asynccontextmanager
from typing import AsyncGenerator, AsyncIterator

from alembic import command
from alembic.config import Config
from sqlalchemy.ext.asyncio import AsyncSession

from src.core.config import settings
from src.db.base import SessionLocal


@asynccontextmanager
async def db_session() -> AsyncGenerator[AsyncSession, None]:
    db = SessionLocal()
    try:
        yield db
        await db.commit()
    except Exception:
        await db.rollback()
        raise
    finally:
        await db.close()


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    async with db_session() as db:
        yield db


async def run_migrations() -> None:
    alembic_cfg = Config("alembic.ini")
    await asyncio.to_thread(command.upgrade, alembic_cfg, "head")


async def create_tables() -> None:
    if not settings.DEBUG:
        await run_migrations()