import aioredis
from aiogram import types
from aiogram.dispatcher.handler import CancelHandler, current_handler
from aiogram.dispatcher.middlewares import BaseMiddleware


def anti_flood(key=1):
    def decorator(func):
        setattr(func, 'throttling_key', key)
        return func

    return decorator


class ThrottlingMiddleware(BaseMiddleware):
    def __init__(self):
        super(ThrottlingMiddleware, self).__init__()

    async def on_process_message(self, answer: types.Message, data: dict):
        handler = current_handler.get()
        throttling_key = getattr(handler, 'throttling_key', None)
        redis_rate = await aioredis.create_redis_pool(
            'redis://server:port',
            password='secret'
        )
        if not await redis_rate.exists(answer.chat.id):
            await redis_rate.setex(
                answer.chat.id,
                throttling_key,
                value=''
            )
            redis_rate.close()
            await redis_rate.wait_closed()
            return
        else:
            redis_rate.close()
            await redis_rate.wait_closed()
            await answer.answer('test')
            raise CancelHandler