"""DB Connector.

Manages interactions with database schema.
"""

import functools
import logging
import time
from contextlib import asynccontextmanager
from contextvars import ContextVar

from sqlalchemy import exc, pool
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import declarative_base

from product_staging import config, exceptions

logger = logging.getLogger(__name__)


def get_engine():
    """Create engine based on config settings."""
    if config.POOL_CLASS == pool.QueuePool:
        _db_engine = create_async_engine(
            config.DB_URL,
            pool_size=config.POOL_SIZE,
            max_overflow=config.POOL_MAX_OVERFLOW,
            pool_recycle=config.POOL_RECYCLE_MS,
            pool_pre_ping=config.POOL_PRE_PING,
        )
        return _db_engine

    return create_async_engine(config.DB_URL, poolclass=config.POOL_CLASS)


db_engine = get_engine()
db_session_maker = async_sessionmaker(db_engine, expire_on_commit=False)
_current_session: ContextVar[AsyncSession | None] = ContextVar(
    "current_session", default=None
)

base_model = declarative_base()


@asynccontextmanager
async def db_session():
    """Provide a transactional scope around a series of operations.

    Taken from http://docs.sqlalchemy.org/en/latest/orm/session_basics.html
    This handles rollback and closing of session, so there is no need
    to do that throughout the code.

    Usage:
        with db_session() as session:
            session.execute(query)
    """
    existing = _current_session.get()
    if existing:
        yield existing
        return

    async with db_session_maker() as session:
        token = _current_session.set(session)
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()
            _current_session.reset(token)


def db_session_wrap(func):
    """Wrap shared DB Session.

    Creates a new session if one isn't passed in. This lets us share/pass-in
    a common session across multiple functions, making them all transactional.
    """

    async def wrapper(*args, **kwargs):
        session = kwargs.pop("session", None)
        if session:
            return await func(*args, session=session, **kwargs)
        else:
            async with db_session() as session:
                return await func(*args, session=session, **kwargs)

    return wrapper


def wrap_db_errors(function):
    """Decorate the given function with logic to handle SQLAlchemy errors.

    If a SQLAlchemy exception is thrown, it will be caught and logged and the
    function will return a fatal response.

    Args:
        function (func): the function to decorate

    Returns:
        func: function decorated with error-handling logic
    """

    @functools.wraps(function)
    async def call_function_with_error_handling(*args, **kwargs):
        try:
            function_return = await function(*args, **kwargs)
        except exc.SQLAlchemyError as exception:
            raise exceptions.SQLException(exception)

        return function_return

    return call_function_with_error_handling


def db_session_wrap_retry_on_deadlock(func):
    """Same as db_session_wrap but will retry the function on OperationalError."""

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        async def retry(func, session, *args, **kwargs):
            max_attempts = 3
            for retry_count in range(1, max_attempts + 1):
                try:
                    return await func(*args, session=session, **kwargs)
                except exc.OperationalError:
                    await session.rollback()
                    if retry_count >= max_attempts:
                        raise
                    time.sleep(retry_count)

        session = kwargs.pop("session", None)
        if session:
            return await retry(func, session, *args, **kwargs)
        else:
            async with db_session() as session:
                return await retry(func, session, *args, **kwargs)

    return wrapper