import os
from typing import AsyncGenerator, Optional

from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import DeclarativeBase

engine: Optional[AsyncEngine] = None
async_session: Optional[async_sessionmaker[AsyncSession]] = None


class Base(DeclarativeBase):
    pass


def init_db_engine() -> None:
    global engine, async_session

    database_url = os.getenv("DATABASE_URL")

    if not database_url:
        raise RuntimeError("DATABASE_URL not found in environment")

    engine = create_async_engine(
        database_url,
        echo=True,
        pool_pre_ping=True,
    )

    async_session = async_sessionmaker(
        bind=engine,
        expire_on_commit=False,
    )


def get_engine() -> AsyncEngine:
    if engine is None:
        raise RuntimeError(
            "DB engine not initialized. Call init_db_engine() first."
        )
    return engine


async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
    if async_session is None:
        raise RuntimeError("DB not initialized. Call init_db_engine() first.")

    async with async_session() as session:
        yield session