package cn.sikey.mcdisk.aspect;

import cn.hutool.http.HttpStatus;
import cn.sikey.framework.common.exception.ServiceException;
import cn.sikey.framework.common.util.servlet.ServletUtils;
import cn.sikey.mcdisk.enums.RateLimit;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Component;
import org.springframework.util.DigestUtils;
import org.springframework.web.multipart.MultipartFile;

import java.lang.reflect.Method;
import java.nio.charset.Charset;
import java.util.Collections;


/**
 * @Author: nelson
 * @Date: 2025/5/19
 * @Description: 限流
 */
@Aspect
@Component
@Slf4j
public class RedisLimitAspect {
    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    @Around("@annotation(rateLimit)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        String parsedKey = buildLimitKey(joinPoint, rateLimit.type());
        log.info("[限流]key值：{}", parsedKey);

        // 使用Lua脚本保证原子操作
        String luaScript = "local current = redis.call('incr', KEYS[1])\n" +
                "if current == 1 then\n" +
                "    redis.call('expire', KEYS[1], ARGV[1])\n" +
                "end\n" +
                "return current";

        Long count = redisTemplate.execute(
                new DefaultRedisScript<>(luaScript, Long.class),
                Collections.singletonList(parsedKey),
                rateLimit.expire()
        );

        if (count > rateLimit.limit()) {
            throw new ServiceException(HttpStatus.HTTP_INTERNAL_ERROR, "请求过于频繁，请稍后再试");
        }

        return joinPoint.proceed();
    }

    private String buildLimitKey(ProceedingJoinPoint joinPoint, RateLimit.LimitType type) {
        StringBuilder keyBuilder = new StringBuilder("rate_limit:");

        // 处理通用标识（用户、IP等）
        switch (type) {
            case USER:
                keyBuilder.append(getUserId());
                break;
            case IP:
                keyBuilder.append(getClientIP());
                break;
            case GLOBAL:
                break;
        }

        Signature signature = joinPoint.getSignature();
        MethodSignature methodSignature = (MethodSignature) signature;
        Method method = methodSignature.getMethod();
        //类名
        keyBuilder.append(joinPoint.getTarget().getClass().getName())
                //方法名
                .append(method.getName());
        for (Object o : joinPoint.getArgs()) {
            if (o instanceof MultipartFile) {
                //附件
                keyBuilder.append(((MultipartFile) o).getOriginalFilename());
            } else {
                keyBuilder.append(o.toString());
            }
        }
        String key = DigestUtils.md5DigestAsHex(keyBuilder.toString().getBytes(Charset.defaultCharset()));
        log.info("[限流]参数：{}", keyBuilder);
        return key;
    }

    private String getUserId() {
        // 从安全上下文获取用户ID（需要根据具体系统实现）
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        return authentication != null ? authentication.getName() : "anonymous";
    }

    private String getClientIP() {
        return ServletUtils.getClientIP();
    }
}

