from typing import Optional, cast

from flask import Flask
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, scoped_session, sessionmaker

from tofu_api.common.config import Config

__all__ = [
    'SQLAlchemy',
]


class SQLAlchemy:
    """
    Wrapper class for integrating SQLAlchemy into the application as a Flask extension.
    """
    _engine: Optional[Engine] = None
    _scoped_session: Optional[scoped_session] = None

    def init_database(self, app: Flask):
        """
        Initializes the SQLAlchemy engine.
        """
        self._engine = self._create_engine(app.config)
        self._scoped_session = self._create_scoped_session()

        @app.teardown_appcontext
        def shutdown_session(_exception=None):
            self._scoped_session.remove()

    @staticmethod
    def _create_engine(config: Config) -> Engine:
        """
        Create the database engine using the app configuration.
        """
        return create_engine(
            config.sqlalchemy_database_uri,
            echo=config.sqlalchemy_echo,
        )

    def _create_scoped_session(self) -> scoped_session:
        """
        Create a scoped session.
        """
        return scoped_session(
            sessionmaker(
                autocommit=False,
                autoflush=False,
                bind=self._engine,
            )
        )

    @property
    def engine(self) -> Engine:
        """
        Database engine.
        """
        assert self._engine, 'Engine not ready yet.'
        return self._engine

    @property
    def session(self) -> Session:
        """
        Scoped database session.
        """
        assert self._scoped_session, 'Session not ready yet.'
        # For all further purposes, the scoped session should be treated like a regular Session object.
        # Use cast() so we can use Session as the type annotation.
        return cast(Session, self._scoped_session)