import aioredis
from aiogram import types, Dispatcher
from aiogram.dispatcher.handler import CancelHandler, current_handler
from aiogram.dispatcher.middlewares import BaseMiddleware
from aiogram.utils.exceptions import Throttled

from config import DEFUALT_RATE_LIMIT, MAX_STANDART_FLOOD_TIME
from language import language


def anti_flood(limit: int, key=None):
    def decorator(func):
        setattr(func, 'throttling_limit', limit)
        return func

    return decorator


# noinspection PyAttributeOutsideInit
class ThrottlingMiddleware(BaseMiddleware):

    def __init__(self, limit=DEFUALT_RATE_LIMIT, key_prefix='antiflood_'):
        self.rate_limit = limit
        self.prefix = key_prefix
        super(ThrottlingMiddleware, self).__init__()

    async def on_process_message(self, message: types.Message, data: dict):
        handler = current_handler.get()
        self.lang = data['lang']

        if handler:
            self.limit = getattr(handler, 'throttling_limit', self.rate_limit)
            self.key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")

        else:
            self.limit = self.rate_limit
            self.key = f"{self.prefix}_message"

        if self.limit <= MAX_STANDART_FLOOD_TIME:
            await message.reply(self.limit)
            await self.standart_flood(message)

        else:
            await self.redis_flood(message)

    async def standart_flood(self, message: types.Message):
        dispatcher = Dispatcher.get_current()

        try:
            await dispatcher.throttle(self.key, rate=self.limit)

        except Throttled as t:
            await self.message_throttled(message, t)
            raise CancelHandler

    async def message_throttled(self, message: types.Message, throttled: Throttled):

        if throttled.exceeded_count == 2:
            await message.reply(language['throttled_request'][self.lang])

    async def redis_flood(self, message: types.Message):

        handler = current_handler.get()
        throttling_limit = getattr(handler, 'throttling_limit', self.limit)

        redis = await aioredis.create_redis_pool(
            'redis://server:port',
            password='password'
        )

        if not await redis.exists(message.chat.id):
            await redis.setex(
                message.chat.id,
                throttling_limit,
                value=handler.__name__
            )

            redis.close()
            await redis.wait_closed()
            return

        elif handler.__name__ == await redis.get(message.chat.id, encoding='utf-8'):

            await message.reply(f"{language['throttled_request'][self.lang]}\n"
                                f"{language['throttled_time'][self.lang].format(await redis.ttl(message.chat.id))}")

            redis.close()
            await redis.wait_closed()
            raise CancelHandler