package com.qmth.boot.api.interceptor.impl;

import com.qmth.boot.api.config.ApiConfig;
import com.qmth.boot.api.config.ApiConfigService;
import com.qmth.boot.api.config.ApiProperties;
import com.qmth.boot.api.constant.ApiConstant;
import com.qmth.boot.api.exception.DefaultExceptionEnum;
import com.qmth.boot.api.interceptor.AbstractInterceptor;
import com.qmth.boot.api.utils.RequestUtil;
import com.qmth.boot.core.rateLimit.entity.RateLimitRule;
import com.qmth.boot.core.rateLimit.entity.RateLimitTarget;
import com.qmth.boot.core.rateLimit.service.RateLimitService;
import com.qmth.boot.core.rateLimit.service.RateLimiter;
import com.qmth.boot.core.security.model.AccessEntity;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.ArrayDeque;
import java.util.Deque;

@Component
public class RateLimitInterceptor extends AbstractInterceptor {

    @Resource
    private ApiProperties apiProperties;

    @Resource
    private ApiConfigService apiConfigService;

    @Resource
    private RateLimitService rateLimitService;

    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        ApiConfig ac = getApiConfig(apiProperties, apiConfigService, request, handler);
        String endpoint = request.getServletPath();
        String device = getDeviceIdOrIpAddress(request);
        String user = getUserIdentity(request);
        // 已通过限流验证堆栈
        Deque<RateLimiter> stack = new ArrayDeque<>();
        // 所有限流规则需要全部匹配才通过
        for (RateLimitRule rule : ac.getRateLimit()) {
            String target = null;
            if (rule.getTarget() == RateLimitTarget.DEVICE) {
                target = device;
            } else if (rule.getTarget() == RateLimitTarget.USER) {
                target = user;
            }
            RateLimiter limiter = rateLimitService.getRateLimiter(endpoint, target, rule);
            if (limiter.acquire()) {
                //限流验证成功则进入堆栈
                stack.push(limiter);
            } else {
                //限流验证失败时，已通过的堆栈需要释放
                while (!stack.isEmpty()) {
                    stack.pop().release();
                }
                log.warn("Rate limited: endpoint={}, rule={}, target={}", endpoint, rule.toString(),
                        StringUtils.trimToEmpty(target));
                throw DefaultExceptionEnum.RATE_LIMITED.exception();
            }
        }
        request.setAttribute(ATTRIBUTE_RATE_LIMITER_STACK, stack);
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler,
            Exception ex) {
        //接口请求完成时，已通过的堆栈需要释放
        Deque<RateLimiter> stack = RequestUtil.getAttribute(request, ApiConstant.ATTRIBUTE_RATE_LIMITER_STACK);
        if (stack != null) {
            while (!stack.isEmpty()) {
                stack.pop().release();
            }
        }
    }

    private String getDeviceIdOrIpAddress(HttpServletRequest request) {
        // deviceId不存在时取IP地址作为设备标识
        String device = StringUtils.trimToNull(request.getHeader(HEADER_DEVICE_ID));
        if (device == null) {
            device = RequestUtil.getIpAddress(request);
        }
        return device;
    }

    private String getUserIdentity(HttpServletRequest request) {
        // 有鉴权信息的情况下，使用签名中的identity作为访问用户标识
        AccessEntity accessEntity = RequestUtil.getAttribute(request, ApiConstant.ATTRIBUTE_ACCESS_ENTITY);
        if (accessEntity != null) {
            return accessEntity.getIdentity();
        } else {
            return StringUtils.EMPTY;
        }
    }

}
