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 