
package cn.sikey.mcdisk.util;

import cn.sikey.framework.common.exception.ServiceException;
import cn.sikey.mcdisk.controller.app.wristwatch.vo.TriggerSmsVerificationReqVO;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.time.temporal.ChronoUnit;
import java.util.Collections;

@Component
public class SmsRateLimiterUtil {

    @Resource
    private RedisTemplate<String, String> redisTemplate;

    private RedisScript<Long> rateLimitScript;

    private static final String RATE_LIMIT_SCRIPT =
            "local key = KEYS[1]\n" +
                    "local maxAttempts = tonumber(ARGV[1])\n" +
                    "local expireSeconds = tonumber(ARGV[2])\n" +
                    "local current = redis.call('GET', key)\n" +
                    "local count = 1\n" +
                    "if current then\n" +
                    "    count = tonumber(current) + 1\n" +
                    "    redis.call('SET', key, count, 'EX', expireSeconds)\n" +
                    "else\n" +
                    "    redis.call('SETEX', key, expireSeconds, count)\n" +
                    "end\n" +
                    "if count > maxAttempts then\n" +
                    "    return -1\n" +
                    "else\n" +
                    "    return count\n" +
                    "end";

    @PostConstruct
    public void init() {
        this.rateLimitScript = new DefaultRedisScript<>(
                RATE_LIMIT_SCRIPT,
                Long.class
        );
    }

    public TriggerSmsVerificationReqVO checkAndIncrement(String phoneNumber, int maxAttempts) {
        String countKey = "SMS_COUNT:" + phoneNumber;
        LocalDateTime now = LocalDateTime.now();
        LocalDateTime midnight = now.plusDays(1).truncatedTo(ChronoUnit.DAYS);
        long secondsUntilMidnight = now.until(midnight, ChronoUnit.SECONDS);

        Long result = redisTemplate.execute(
                rateLimitScript,
                Collections.singletonList(countKey),
                maxAttempts,
                secondsUntilMidnight
        );

        if (result == null || result == -1) {
            throw new ServiceException(HttpStatus.TOO_MANY_REQUESTS.value(),
                    "短信发送次数超过限制");
        }

        TriggerSmsVerificationReqVO triggerSmsVerificationReqVO = new TriggerSmsVerificationReqVO();
        triggerSmsVerificationReqVO.setSmsCount(result);
        return triggerSmsVerificationReqVO;
    }
}
