import os

from fastapi import Depends
from fastapi_users_db_sqlmodel import SQLModelBaseUserDB
from fastapi_users_db_sqlmodel import SQLModelUserDatabaseAsync
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlmodel import Field


class Base(SQLModelBaseUserDB):
    pass


class User(Base, table=True):
    username: str = Field(unique=True)


database_url = (
    f'postgresql+asyncpg://{os.environ.get("POSTGRES_USER")}:'
    + f'{os.environ.get("POSTGRES_PASSWORD")}@'
    + os.environ.get("POSTGRES_CONTAINER")
    + f':5432/{os.environ.get("POSTGRES_DB")}'
)

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


async def get_session() -> AsyncSession:
    async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    async with async_session() as session:
        yield session


async def create_db_and_tables():
    # TODO: Remove the drop, this is for dev
    async with engine.begin() as conn:
        # await conn.run_sync(Base.metadata.drop_all)
        await conn.run_sync(Base.metadata.create_all)


async def get_user_db(session: AsyncSession = Depends(get_session)):
    yield SQLModelUserDatabaseAsync(session, User)
