import time
from contextlib import asynccontextmanager

from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession

from app.core._config import DATABASE_URL
from app.core.logger import make_log


def _to_async_dsn(url: str) -> str:
    # Convert psycopg2 DSN to asyncpg DSN
    # postgresql+psycopg2://user:pass@host:5432/db -> postgresql+asyncpg://user:pass@host:5432/db
    return url.replace("+psycopg2", "+asyncpg")


# Async engine for PostgreSQL
engine = create_async_engine(
    _to_async_dsn(DATABASE_URL),
    pool_size=10,
    max_overflow=20,
    pool_timeout=30,
    pool_recycle=1800,
    pool_pre_ping=True,
)

AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)


async def wait_db_ready():
    ready = False
    while not ready:
        try:
            async with engine.connect() as conn:
                await conn.execute(text("SELECT 1"))
                ready = True
        except Exception as e:
            make_log("SQL", 'PostgreSQL is not ready yet: ' + str(e), level='debug')
            time.sleep(1)


@asynccontextmanager
async def db_session(auto_commit: bool = False):
    session: AsyncSession = AsyncSessionLocal()
    try:
        yield session
        if auto_commit:
            await session.commit()
    except BaseException as e:
        await session.rollback()
        raise e
    finally:
        await session.close()


def new_session() -> AsyncSession:
    return AsyncSessionLocal()