"""Database Connector for SQLAlchemy."""
from collections.abc import Generator
from contextlib import contextmanager
from contextvars import ContextVar
import os
from typing import Any
from uuid import uuid4

from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session
from sqlalchemy.orm import Session
from sqlalchemy.orm import sessionmaker
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.base import RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp


class Database:
    """Setup and contain our database connection.

    - This is used to be able to setup the database in an uniform way
    while allowing easy testing and session management.
    - Session management is done using ``scoped_session`` with a special scopefunc,
    because we cannot use threading.local(). Contextvar does the right thing
    with respect to asyncio and behaves similar to threading.local().
    - We only store a random string in the contextvar and let scoped session do
    the heavy lifting.
    - This allows us to easily start a new session or get the existing one using
    the scoped_session mechanics.
    """

    def __init__(self, db_url: str, engine_args: dict, session_args: dict) -> None:
        """Initialise the database object with the db url, engine arguments and session arguments.

        Args:
            db_url (str): url of the database
            engine_args (dict): a dictionary containing the engine arguments
            session_args (dict): a dictionary containing the session arguments
        """
        self.request_context: ContextVar[str] = ContextVar('request_context', default='')

        if os.environ.get('test_type') == 'unit':
            self.engine = create_engine('sqlite://')
        else:
            self.engine = create_engine(db_url, **engine_args)
        session_args.update({'class_': Session})
        self.session_factory = sessionmaker(bind=self.engine, **session_args)

        self.scoped_session = scoped_session(self.session_factory, self._scopefunc)

    def _scopefunc(self) -> str | None:
        scope_str = self.request_context.get()
        return scope_str

    @property
    def session(self) -> Session:
        """Return a database session."""
        return self.scoped_session()

    @contextmanager
    def create_database_scope(self, **kwargs: Any) -> Generator['Database', None, None]:
        """Create a new database session (scope).

        - This creates a new database session to handle all the database connection
        from a single scope (request or workflow).
        - This method should typically only been called in request middleware or
        at the start of workflows.
        Args:
            ``**kwargs``: Optional session kw args for this session
        """
        token = self.request_context.set(str(uuid4()))
        self.scoped_session(**kwargs)
        yield self
        self.scoped_session.remove()
        self.request_context.reset(token)


class DBSessionMiddleware(BaseHTTPMiddleware):
    """Creates Database Session Middleware."""

    def __init__(self, app: ASGIApp, database: Database, commit_on_exit: bool = False):
        """Initialise the DBSessionMiddleware.

        Args:
            app (ASGIApp): the ASGI app
            database (Database): The database to be used
            commit_on_exit (bool): a boolean flag to commit on exit
        """
        super().__init__(app)
        self.commit_on_exit = commit_on_exit
        self.database = database

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        """Dispatch overrides method of class BaseHTTPMiddleware.

        Args:
            request(Request): the request object
            call_next(RequestResponseEndpoint): the request response endpoint
        Returns:
            Response: a starlette response object
        """
        with self.database.create_database_scope():
            response = await call_next(request)
        return response