# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
# OpenOLC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         `http://license.coscl.org.cn/MulanPSL2`
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

import logging
from typing import List, Any, Optional

from olc.bean.constant import QUOTA_SCRIPT_NAME, QUOTA_SCRIPT_ROLLBACK_NAME
from olc.bean.match_wrapper import MatchWrapper
from olc.bean.olc_config_rule import FlowPolicy
from olc.limit.cluster.abs_cluster_limiter import AbsClusterLimiter
from olc.limit.cluster.redis.async_redis_client import AsyncOlcRedisManager
from olc.limit.cluster.redis.redis_client import OlcRedisManager
from olc.limit.node.quota_bucket import QuotaBucket
from olc.utils.empty_checker import EmptyChecker
from olc.utils.time_utils import TimeUnit

logger = logging.getLogger(__name__)


class QuotaClusterLimiter(AbsClusterLimiter):
    """
    配额集群限制器
    """

    def __init__(self, wrapper: MatchWrapper[FlowPolicy]):
        super().__init__(wrapper)
        self.policy = wrapper.policy
        self.group = wrapper.tag_group
        self._limit = self.policy.rate_limit
        self._burst_limit = self.policy.burst_limit
        self._remote_key = "olc##quota_cluster_limiter##" + self.name
        # 配额窗口大小，单位秒
        self._window_size = self.policy.time_unit_to_seconds()
        # 初始化降级用的本地令牌桶
        self._init_fallback_bucket()

    # ========== 公共方法 同/异步公用 ==================
    def _get_rollback_script_key(self) -> str:
        return QUOTA_SCRIPT_ROLLBACK_NAME

    def _get_script_key(self) -> str:
        return QUOTA_SCRIPT_NAME

    def _get_remote_key(self) -> List[str]:
        return [self._remote_key]

    def _get_remote_value(self, token_number: int) -> List[Any]:
        # lua 脚本参数顺序
        # 1. burst_limit
        # 2. token_number
        # 3. window_size
        return [self._burst_limit, token_number, self._window_size]

    def _handle_result(self, result: List[Any], token_number: int) -> bool:
        logger.debug("%s get token result: %s", self._remote_key, result)
        if EmptyChecker.is_empty_list(result):
            logger.warning(
                "%s get token result is empty, bypass request",
                self._remote_key
            )
            return True
        return result[0] >= token_number

    def _init_fallback_bucket(self):
        if self._safe_mode == 1:
            rate_limit = self.policy.assign_alg.get_safe_rate_limit()
            if rate_limit is None:
                rate_limit = self.policy.rate_limit
            self._fallback_bucket = QuotaBucket(
                time_interval=1,
                time_unit=TimeUnit.SECOND,
                rate_limit=rate_limit,
            )
        logger.debug(
            "%s fallback bucket set, safe mode is %s",
            self.name, self._safe_mode
        )

    # ============== 同步 处理逻辑================
    def pre_check(self, token_number: int) -> bool:
        return True

    def roll_back(self, token_number: int):
        """
            回滚令牌
        :param token_number:
        :return:
        """
        if token_number <= 0:
            return

        if self.is_exception_status():
            if self._safe_mode == 1:
                self._fallback_bucket.add_token_number(token_number)
                return
        try:
            result = OlcRedisManager.execute_script(self._get_rollback_script_key(), self._get_remote_key(),
                                                    args=[2, token_number, self._burst_limit])
            logger.debug("%s rollback result: %s", self.name, result)
        except Exception as e:
            logger.warning("%s rollback error: %s", self.name, e)

    def dec_token(self, token_number: int):
        """
            强制扣减令牌， 扣减至0
        :param token_number:
        :return:
        """
        if self.is_exception_status():
            if self._safe_mode == 1:
                if token_number > 0:
                    self._fallback_bucket.dec_token_number(token_number)
                if token_number < 0:
                    self._fallback_bucket.add_token_number(abs(token_number))
                return

        if token_number == 0:
            return
        try:
            if token_number > 0:
                result = OlcRedisManager.execute_script(self._get_rollback_script_key(), self._get_remote_key(),
                                               args=[1, token_number, self._burst_limit])
            else:
                # 添加令牌
                result = OlcRedisManager.execute_script(
                    self._get_rollback_script_key(),
                    keys=self._get_remote_key(),
                    args=[2, abs(token_number), self._burst_limit],
                )
            logger.debug("%s dec_token result: %s", self.name, result)
        except Exception as e:
            logger.warning("%s dec_token error: %s", self.name, e)

    # ============== 异步处理逻辑================
    async def async_roll_back(self, token_number: int):
        if token_number <= 0:
            return
        if self.is_exception_status():
            if self._safe_mode == 1:
                self._fallback_bucket.add_token_number(token_number)
            return
        try:
            result = await AsyncOlcRedisManager.execute_script(
                self._get_rollback_script_key(),
                self._get_remote_key(),
                [2, token_number, self._burst_limit],
            )
            logger.debug("%s async rollback result: %s", self.name, result)
        except Exception as e:
            logger.warning("%s async rollback error: %s", self.name, e)

    async def async_dec_token(self, token_number: int):
        if self.is_exception_status():
            if self._safe_mode == 1:
                if token_number > 0:
                    self._fallback_bucket.dec_token_number(token_number)
                if token_number < 0:
                    self._fallback_bucket.add_token_number(abs(token_number))
            return

        if token_number == 0:
            return
        try:
            if token_number > 0:
                args = [1, token_number, self._burst_limit]
            else:
                args = [2, abs(token_number), self._burst_limit]
            result = await AsyncOlcRedisManager.execute_script(
                self._get_rollback_script_key(),
                self._get_remote_key(),
                args,
            )
            logger.debug("%s async dec_token result: %s", self.name, result)
        except Exception as e:
            logger.warning("%s async dec_token error: %s", self.name, e)