package top.tqx.demo_1.aspect;

import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import top.tqx.demo_1.annotation.RateLimit;
import top.tqx.demo_1.util.IpUtil;
import top.tqx.demo_1.util.RedisUtil;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Aspect
@Component
@Slf4j
public class RateLimitAspect {

    @Autowired(required = false)
    private RedisUtil redisUtil;

    private final ConcurrentHashMap<String, RateLimitInfo> localCache = new ConcurrentHashMap<>();

    @Around("@annotation(rateLimit)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes == null) {
            return joinPoint.proceed();
        }

        HttpServletRequest request = attributes.getRequest();

        if (redisUtil != null) {
            try {
                return handleWithRedis(joinPoint, rateLimit, request);
            } catch (Exception e) {
                log.warn("Redis 限流失败，降级使用本地缓存: {}", e.getMessage());
            }
        }

        return handleWithLocalCache(joinPoint, rateLimit, request);
    }

    private Object handleWithRedis(ProceedingJoinPoint joinPoint, RateLimit rateLimit, HttpServletRequest request) throws Throwable {
        String key = generateKey(request, rateLimit);
        Long count = redisUtil.increment(key);
        if (count == 1) {
            redisUtil.expire(key, rateLimit.timeWindow(), TimeUnit.SECONDS);
        }

        if (count > rateLimit.maxRequests()) {
            log.warn("IP: {} 触发限流, 接口: {}", IpUtil.getIpAddress(request), request.getRequestURI());
            throw new RuntimeException("请求过于频繁，请稍后再试");
        }

        return joinPoint.proceed();
    }

    private Object handleWithLocalCache(ProceedingJoinPoint joinPoint, RateLimit rateLimit, HttpServletRequest request) throws Throwable {
        String key = generateKey(request, rateLimit);
        long now = System.currentTimeMillis();

        RateLimitInfo info = localCache.computeIfAbsent(key, k -> new RateLimitInfo());

        if (now - info.windowStart > rateLimit.timeWindow() * 1000) {
            info.windowStart = now;
            info.count = 0;
        }

        info.count++;

        if (info.count > rateLimit.maxRequests()) {
            log.warn("IP: {} 触发限流(本地), 接口: {}", IpUtil.getIpAddress(request), request.getRequestURI());
            throw new RuntimeException("请求过于频繁，请稍后再试");
        }

        return joinPoint.proceed();
    }

    private String generateKey(HttpServletRequest request, RateLimit rateLimit) {
        String prefix = "rate_limit:";
        if ("ip".equals(rateLimit.key())) {
            return prefix + IpUtil.getIpAddress(request) + ":" + request.getRequestURI();
        } else if ("user".equals(rateLimit.key())) {
            String userId = request.getHeader("X-User-Id");
            return prefix + (userId != null ? userId : "anonymous") + ":" + request.getRequestURI();
        }
        return prefix + "global:" + request.getRequestURI();
    }

    static class RateLimitInfo {
        long windowStart;
        long count;
    }
}