# 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.

from olc.bean.match_wrapper import MatchWrapper
from olc.bean.olc_config_rule import FlowPolicy
from olc.limit.limiter import Limit
from olc.limit.node.qps_token_bucket import QpsTokenBucket


class QpsLimiter(Limit):

    def __init__(self, wrapper: MatchWrapper[FlowPolicy]):
        self._wrapper = wrapper
        self.alg = QpsTokenBucket.create(wrapper)

    def get_match_wrapper(self) -> MatchWrapper[FlowPolicy]:
        return self._wrapper

    def try_acquire(self, token_number: int) -> bool:
        return self.alg.try_acquire(token_number)

    def pre_check(self, token_number: int) -> bool:
        return True

    def roll_back(self, token_number: int):
        return self.alg.add_token_number(token_number)

    def dec_token(self, token_number: int):
        if token_number == 0:
            return
        if token_number > 0:
            self.alg.dec_token_number(token_number)
        else:
            self.alg.dec_token_number(abs(token_number))

    async def async_try_acquire(self, token_number: int) -> bool:
        return await self.alg.async_try_acquire(token_number)


    async def async_try_acquire(self, token_number: int) -> bool:
        return await self.alg.async_try_acquire(token_number)