AIyaodianzhushou/backend-java/src/main/java/com/pharmacopoeia/service/DashScopeRateLimiter.java

package com.pharmacopoeia.service;

import com.pharmacopoeia.config.DashScopeRateLimitProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import java.time.LocalDate;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 百炼 API 全局调用限流器（内存计数，重启清零）。
 * 控制后端 → 百炼的总调用量，与用户侧限流独立。
 */
@Service
public class DashScopeRateLimiter {

    private static final Logger log = LoggerFactory.getLogger(DashScopeRateLimiter.class);

    private final DashScopeRateLimitProperties props;
    private final AtomicInteger minuteCount = new AtomicInteger(0);
    private final AtomicInteger hourCount = new AtomicInteger(0);
    private final AtomicInteger dayCount = new AtomicInteger(0);
    private volatile String currentDate = LocalDate.now().toString();
    private volatile int currentHour = java.time.LocalTime.now().getHour();

    public DashScopeRateLimiter(DashScopeRateLimitProperties props) {
        this.props = props;
    }

    /** 尝试获取一次调用许可。返回 true 表示放行，false 表示超限。 */
    public boolean tryAcquire() {
        if (!props.isEnabled()) return true;

        resetIfNeeded();

        if (props.getPerMinute() > 0 && minuteCount.incrementAndGet() > props.getPerMinute()) {
            log.warn("百炼 API 每分钟限流触发: {}/{}", minuteCount.get(), props.getPerMinute());
            return false;
        }
        if (props.getPerHour() > 0 && hourCount.incrementAndGet() > props.getPerHour()) {
            log.warn("百炼 API 每小时限流触发: {}/{}", hourCount.get(), props.getPerHour());
            return false;
        }
        if (props.getPerDay() > 0 && dayCount.incrementAndGet() > props.getPerDay()) {
            log.warn("百炼 API 每日限流触发: {}/{}", dayCount.get(), props.getPerDay());
            return false;
        }
        return true;
    }

    private void resetIfNeeded() {
        String today = LocalDate.now().toString();
        int nowHour = java.time.LocalTime.now().getHour();
        if (!currentDate.equals(today)) {
            dayCount.set(0);
            currentDate = today;
        }
        if (currentHour != nowHour) {
            hourCount.set(0);
            currentHour = nowHour;
        }
    }
}