from aiogram import BaseMiddleware
from aiogram.types import TelegramObject, Message, CallbackQuery
from core.config import config
from core.emojis import p_em
from core.utils import check_user_membership, get_join_required_markup
from database.database import AsyncSessionLocal

class MandatoryJoinMiddleware(BaseMiddleware):
    async def __call__(self, handler, event: TelegramObject, data: dict):
        user = None
        bot = data.get('bot')
        
        if isinstance(event, Message):
            user = event.from_user
        elif isinstance(event, CallbackQuery):
            user = event.from_user
            
        if user and bot:
            # Allow admins to bypass
            if user.id in config.admin_ids:
                return await handler(event, data)
                
            # Allow check_membership callback to bypass so it can trigger the actual handler
            if isinstance(event, CallbackQuery) and event.data == "check_membership":
                return await handler(event, data)
                
            async with AsyncSessionLocal() as session:
                is_member, unjoined_links = await check_user_membership(bot, user.id, session)
                
            if not is_member:
                markup = get_join_required_markup(unjoined_links)
                text = p_em(
                    "🛑 <b>دسترسی محدود است!</b>\n\n"
                    "برای استفاده از امکانات ربات، باید در کانال(های) زیر عضو باشید:\n\n"
                    "👇 لطفاً پس از عضویت، روی <b>بررسی عضویت</b> کلیک کنید:"
                )
                
                if isinstance(event, Message):
                    await event.answer(text, reply_markup=markup)
                elif isinstance(event, CallbackQuery):
                    await event.message.edit_text(text, reply_markup=markup)
                    await event.answer()
                return # Block execution
                
        return await handler(event, data)