# 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 asyncio
import sys
import threading
import time

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 QpsTokenBucket(TokenService):

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

    def __init__(self, time_interval: int, time_unit: TimeUnit, rate_limit: int, buster_limit: int,
                 wait_time_ms: int = 0, need_refresh: bool = True):
        """
            初始化令牌桶
        :param time_interval: 时间间隔
        :param time_unit: 时间单位
        :param rate_limit: 速率
        :param buster_limit: 桶大小
        :param wait_time_ms: 等待时间
        :param need_refresh: 是否需要刷新
        """
        self._buster_limit = buster_limit
        # token_interval_ms 参数一个令牌的时间间隔，单位毫秒
        if rate_limit <= 0:
            self._token_interval_ms = sys.maxsize
        else:
            self._token_interval_ms = time_unit.to_millisecond(time_interval) / rate_limit

        self._time_interval = time_interval
        self._time_unit = time_unit
        # 使用 threading.Lock 而非 asyncio.Lock：锁内全部为纯 CPU 计算（微秒级），
        # 不会阻塞事件循环；同时保证同步 try_acquire 与异步 async_try_acquire 共享同一把锁
        self._lock = threading.Lock()
        self._wait_time_ms = wait_time_ms
        self._last_time = AtomicNum(TimeUtils.get_milliseconds())
        self._need_refresh = need_refresh
        if self._need_refresh:
            self._token_number = AtomicNum(float(buster_limit))
        else:
            self._token_number = AtomicNum(float(0))
        self._last_token_number = self._token_number.get()

    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:
            add_tokens = (current_time - last_time) / self._token_interval_ms
            if add_tokens > 0 and self._last_time.compare_and_set(last_time, current_time):
                self._token_number.add_and_limit(add_tokens, self._buster_limit)

    def _acquire(self, num: int) -> bool:
        """
            申请令牌
        :param num: 令牌数量
        :return:
        """
        if self._token_number.get() >= num:
            self._last_token_number = self._token_number.get()
            self._token_number.dec_and_limit(num, 0)
            return True
        return False

    def _calculate_token_wait_time(self, num: int) -> int:
        """计算获取指定数量令牌需要等待的毫秒数（纯计算，不修改任何状态）"""
        current_tokens = self._token_number.get()
        if current_tokens >= num:
            return 0
        needed_tokens = num - current_tokens
        return int(needed_tokens * self._token_interval_ms)

    def _roll_back(self):
        """回滚令牌"""
        current_tokens = self._token_number.get()
        self._token_number.set(self._last_token_number)
        rollback_time = int(
            (current_tokens - self._last_token_number)
            * self._token_interval_ms
        )
        self._last_time.set(self._last_time.get() - rollback_time)

    def _try_acquire_locked(self, token_number: int) -> int:
        """锁内执行令牌获取逻辑，返回需等待的毫秒数；-1 表示获取失败"""
        self._refresh_token()

        if self._wait_time_ms <= 0:
            return 0 if self._acquire(token_number) else -1

        current_time = TimeUtils.get_milliseconds()
        # 计算需要等待token的时间
        wait_time_ms = self._calculate_token_wait_time(token_number)

        # 计算真实需要等待时间, 上个令牌已经占用了部分时间,所有需要计算真实等待时间
        last_time = self._last_time.get()
        real_wait_time = last_time + wait_time_ms - current_time

        if real_wait_time > self._wait_time_ms:
            return -1

        # 记录回滚快照
        self._last_token_number = self._token_number.get()
        # 统一扣减令牌（无论是否需要等待）
        self._token_number.dec_and_limit(token_number, 0)
        # 需要等待时，预推 last_time 以抵消等待期间自动补充的令牌
        if wait_time_ms > 0:
            self._last_time.add(wait_time_ms)
        return max(0, real_wait_time)

    def try_acquire(self, token_number: int = 1) -> bool:
        with self._lock:
            wait_time_ms = self._try_acquire_locked(token_number)

        if wait_time_ms < 0:
            return False
        if wait_time_ms > 0:
            try:
                time.sleep(wait_time_ms / 1000.0)
            except InterruptedError | KeyboardInterrupt:
                # 中途打断什么都不做
                pass
        return True

    async def async_try_acquire(self, num: int = 1) -> bool:
        with self._lock:
            wait_time_ms = self._try_acquire_locked(num)

        if wait_time_ms < 0:
            return False
        if wait_time_ms > 0:
            try:
                await asyncio.sleep(wait_time_ms / 1000.0)
            except asyncio.CancelledError:
                # 协程被取消时直接跳过
                pass
        return True

    def pre_check(self, num: int = 1) -> bool:
        """预检查是否能在 wait_time_ms 内获取到指定数量的令牌"""
        self._refresh_token()
        current_tokens = self._token_number.get()
        if current_tokens >= num:
            return True
        if self._wait_time_ms <= 0:
            return False
        # 计算令牌缺口对应的等待时间
        wait_time_ms = self._calculate_token_wait_time(num)
        # 结合 _last_time 排队 backlog 计算真实等待时间
        current_time = TimeUtils.get_milliseconds()
        real_wait_time = self._last_time.get() + wait_time_ms - current_time
        return real_wait_time <= self._wait_time_ms

    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._buster_limit))

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

    def dec_token_number(self, num: int):
        self._token_number.dec_and_limit(num, 0)

    def update_rate(self, rate_limit: int):
        if rate_limit <= 0:
            self._token_interval_ms = sys.maxsize
        else:
            # 计算新令牌间隔
            new_interval = float(self._time_unit.to_millisecond(self._time_interval)) / rate_limit
            if self._token_interval_ms != sys.maxsize:
                ratio = self._token_interval_ms / new_interval
                # 按比例调整
                if ratio != 1 and ratio > 0:
                    current_token = self._token_number.get()
                    new_token = int(current_token * ratio)
                    self._token_number.set(min(new_token, self._buster_limit))
            self._token_interval_ms = new_interval