from typing import Any, Awaitable, Callable, Dict
from aiogram import BaseMiddleware, Bot
from aiogram.types import TelegramObject
from sqlalchemy import select
from loguru import logger
from database import async_session_, Worker


class SessionMiddleware(BaseMiddleware):
    sessions_in_memory_db = set()

    async def __call__(
            self,
            handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
            event: TelegramObject,
            data: Dict[str, Any]
    ) -> Any:
        logger.info("Session check")
        event_data = event.message or event.callback_query
        user = event_data.from_user.id
        if user not in self.sessions_in_memory_db:
            async with async_session_() as session:
                result = await session.execute(select(Worker).where(Worker.telegram_id == event_data.from_user.id))
                user_in_db = result.scalars().first()
            if not user_in_db:
                data["through_registration"] = True
            else:
                self.sessions_in_memory_db.add(user)
        return await handler(event, data)