package com.reading.platform.common.util;

import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

/**
 * Redis 限流工具类
 * 基于滑动窗口算法实现限流
 */
@Slf4j
@Component
public class RateLimiter {

    private final RedisTemplate<String, Object> redisTemplate;

    public RateLimiter(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    /**
     * 尝试获取请求许可
     *
     * @param key        限流键
     * @param maxRequests 时间窗口内最大请求数
     * @param timeWindow  时间窗口大小
     * @param timeUnit    时间单位
     * @return true-允许请求，false-拒绝请求
     */
    public boolean tryAcquire(String key, long maxRequests, long timeWindow, TimeUnit timeUnit) {
        long currentTime = System.currentTimeMillis();
        long windowStart = currentTime - timeUnit.toMillis(timeWindow);

        String redisKey = "rate_limit:" + key;

        try {
            // 使用 Redis 事务保证原子性
            Boolean result = redisTemplate.execute(connection -> {
                try {
                    // 移除窗口外的请求记录
                    connection.zSetCommands().zRemRangeByScore(
                            redisKey.getBytes(),
                            0,
                            windowStart
                    );

                    // 获取当前窗口内的请求数
                    Long count = connection.zSetCommands().zCard(redisKey.getBytes());

                    if (count != null && count >= maxRequests) {
                        // 超过限流阈值
                        return false;
                    }

                    // 添加当前请求记录
                    connection.zSetCommands().zAdd(
                            redisKey.getBytes(),
                            currentTime,
                            String.valueOf(currentTime).getBytes()
                    );

                    // 设置过期时间
                    connection.expire(redisKey.getBytes(), timeWindow);

                    return true;
                } catch (Exception e) {
                    log.error("限流检查失败，key: {}", key, e);
                    // 异常情况放行
                    return true;
                }
            }, false);

            return Boolean.TRUE.equals(result);
        } catch (Exception e) {
            log.error("限流检查异常，key: {}", key, e);
            // 异常情况放行
            return true;
        }
    }

    /**
     * 获取剩余请求次数
     *
     * @param key        限流键
     * @param maxRequests 时间窗口内最大请求数
     * @param timeWindow  时间窗口大小
     * @param timeUnit    时间单位
     * @return 剩余请求次数
     */
    public long getRemainingRequests(String key, long maxRequests, long timeWindow, TimeUnit timeUnit) {
        long currentTime = System.currentTimeMillis();
        long windowStart = currentTime - timeUnit.toMillis(timeWindow);

        String redisKey = "rate_limit:" + key;

        try {
            Long count = redisTemplate.execute(connection -> {
                // 先清理过期数据
                connection.zSetCommands().zRemRangeByScore(
                        redisKey.getBytes(),
                        0,
                        windowStart
                );
                // 获取当前窗口内的请求数
                return connection.zSetCommands().zCard(redisKey.getBytes());
            }, false);

            if (count == null) {
                return maxRequests;
            }

            return Math.max(0, maxRequests - count);
        } catch (Exception e) {
            log.error("获取剩余请求次数失败，key: {}", key, e);
            return maxRequests;
        }
    }

    /**
     * 重置限流计数器
     *
     * @param key 限流键
     */
    public void reset(String key) {
        String redisKey = "rate_limit:" + key;
        try {
            redisTemplate.delete(redisKey);
            log.debug("限流计数器已重置，key: {}", key);
        } catch (Exception e) {
            log.error("重置限流计数器失败，key: {}", key, e);
        }
    }
}
