# 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 olc.bean.match_wrapper import MatchWrapper
from olc.bean.olc_config_rule import FlowPolicy
from olc.limit.token_service import TokenService
from olc.utils.atomic_num import AtomicNum
from olc.utils.time_utils import TimeUnit, TimeUtils


class QuotaBucket(TokenService):

    @staticmethod
    def create(wrapper: MatchWrapper[FlowPolicy]) -> "QuotaBucket":
        policy = wrapper.policy
        return QuotaBucket(
            policy.time_unit_to_seconds(), TimeUnit.SECOND, policy.rate_limit
        )

    def __init__(
            self,
            time_interval: int,
            time_unit: TimeUnit,
            rate_limit: int,
            need_refresh: bool = True,
    ) -> None:
        """
            固定窗口算法
        :param time_interval: 时间窗长度
        :param time_unit:    时间窗单位
        :param rate_limit:  速率
        :param need_refresh: 是否需要刷新令牌，不刷新的话作为远程令牌本地存储
        """
        self._time_interval = time_interval
        self._time_unit = time_unit
        self._rate_limit = rate_limit
        self._need_refresh = need_refresh
        self._time_interval_ms = TimeUnit.SECOND.to_millisecond(time_interval)
        self._token_number = AtomicNum(rate_limit)
        self._last_time = AtomicNum(TimeUtils.get_milliseconds())

    def _refresh_token(self):
        if not self._need_refresh:
            return
        current_time = TimeUtils.get_milliseconds()
        last_time = self._last_time.get()
        if current_time > last_time + self._time_interval_ms:
            if self._last_time.compare_and_set(last_time, current_time):
                self._token_number.get_and_set(self._rate_limit)

    def pre_check(self, num: int = 1) -> bool:
        self._refresh_token()
        return self._token_number.get() >= num

    def try_acquire(self, num: int = 1) -> bool:
        self._refresh_token()
        return self._token_number.try_decrement(num)

    def get_token_number(self):
        return self._token_number.get()

    def set_token_number(self, token_number: int):
        if token_number > 0:
            self._token_number.set(min(token_number, self._rate_limit))

    def add_token_number(self, num: int):
        self._token_number.add_and_limit(num, self._rate_limit)

    def dec_token_number(self, num: int):
        self._token_number.decrement_to_zero(num)

    def update_rate(self, rate_limit: int):
        """
            更新速率
        :param rate_limit:
        :return:
        """
        if rate_limit <= 0:
            self._rate_limit = 0
            self._token_number.set(0)
            return
        # rate_limit> 0  self._rate_limit<0
        if self._rate_limit <= 0:
            self._rate_limit = rate_limit
            self._token_number.set(rate_limit)
            return
        # rate_limit> 0  self._rate_limit > 0
        ratio = rate_limit / float(self._rate_limit)# 确保 ratio 至少为 1，避免令牌数归零
        self._rate_limit = rate_limit
        self._token_number.multiply(ratio)