from typing import Union

from fastapi import FastAPI
from pydantic import BaseSettings

from ...utils import get_settings

try:
    from sqlalchemy.engine import Engine
    from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
    from sqlmodel import Session, SQLModel, create_engine
    from sqlmodel.ext.asyncio.session import AsyncSession

except ImportError:
    raise RuntimeError(
        "SQLModel is not installed. Please install it with `pip install sqlmodel pyhumps`"
    )


class Database:
    """
    A class to wrap the sqlalchemy engine and open a connection session to the db.
    """

    def __init__(self, engine: Union[Engine, AsyncEngine], is_async: bool = False):
        self.engine = engine
        self.is_async = is_async

    def open(self) -> Union[Session, AsyncSession]:
        if self.is_async:
            return AsyncSession(self.engine)
        else:
            return Session(self.engine)


def setup(app: FastAPI, settings: BaseSettings = None) -> Database:
    """
    Install the sqlmodel plugin to the app.
    This will attach 1 attribute to `app.state` i.e:

    * `db` - `popol.sqlmodel.Database` instance object to open db connection.

    Args:
        app: FastAPI app.
        settings: The settings (can be pydantic.BaseSettings).

    Returns:
        Database: The database.
    """

    settings = get_settings(app, settings)
    prefix = "SQLALCHEMY_"
    db_uri = getattr(settings, f"{prefix}DATABASE_URI", None)
    if not db_uri:
        raise RuntimeError(f"{prefix}DATABASE_URI is not set")

    async_mode = getattr(settings, f"{prefix}ASYNC_MODE", False)
    options = getattr(settings, f"{prefix}OPTIONS", {})
    if async_mode:
        engine = create_async_engine(db_uri, **options)
    else:
        engine = create_engine(db_uri, **options)

    db = Database(engine, async_mode)
    app.state.db = db

    async def startup():
        # reference: https://github.com/tiangolo/sqlmodel/issues/54#issue-981884262
        if async_mode:
            async with engine.begin() as conn:
                await conn.run_sync(SQLModel.metadata.create_all)
        else:
            SQLModel.metadata.create_all(engine)

    app.add_event_handler("startup", startup)
    return db

sqlalchemy.ext.asyncio.create_async_engine sqlmodel.SQLModel.metadata.create_all sqlmodel.Session sqlmodel.create_engine sqlmodel.ext.asyncio.session.AsyncSession
[((1723, 1761), 'sqlalchemy.ext.asyncio.create_async_engine', 'create_async_engine', (['db_uri'], {}), '(db_uri, **options)\n', (1742, 1761), False, 'from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine\n'), ((1789, 1821), 'sqlmodel.create_engine', 'create_engine', (['db_uri'], {}), '(db_uri, **options)\n', (1802, 1821), False, 'from sqlmodel import Session, SQLModel, create_engine\n'), ((862, 887), 'sqlmodel.ext.asyncio.session.AsyncSession', 'AsyncSession', (['self.engine'], {}), '(self.engine)\n', (874, 887), False, 'from sqlmodel.ext.asyncio.session import AsyncSession\n'), ((921, 941), 'sqlmodel.Session', 'Session', (['self.engine'], {}), '(self.engine)\n', (928, 941), False, 'from sqlmodel import Session, SQLModel, create_engine\n'), ((2155, 2191), 'sqlmodel.SQLModel.metadata.create_all', 'SQLModel.metadata.create_all', (['engine'], {}), '(engine)\n', (2183, 2191), False, 'from sqlmodel import Session, SQLModel, create_engine\n')]