from telebot.handler_backends import BaseMiddleware
from telebot.types import Message, CallbackQuery

from sqlalchemy.orm import Session

from ..database.models import User


class DatabaseMiddleware (BaseMiddleware):
    def __init__(self, engine):
        super().__init__()
        self.engine = engine
        self.update_types = ["message", "callback_query"]

    def pre_process(self, obj: [Message, CallbackQuery], data: dict):
        session = Session(self.engine)
        user = session.get(User, obj.from_user.id)
        if user is None:
            user = User(id=obj.from_user.id, username=obj.from_user.username)
            session.add(user)
            session.commit()
        data["db"] = session
        data["user"] = user

    def post_process(self, message, data: dict, exception: BaseException):
        if "db" in data:
            session: Session = data["db"]
            session.rollback()
            session.close()