package com.qmth.boot.core.rateLimit.service.impl;

import com.qmth.boot.core.rateLimit.entity.RateLimitRule;
import com.qmth.boot.core.rateLimit.service.RateLimitService;
import com.qmth.boot.core.rateLimit.service.RateLimiter;

import java.util.HashMap;
import java.util.Map;

public class MemoryRateLimitService implements RateLimitService {

    private static final String TARGET_ALL_KEY = "*";

    private final Map<String, Map<String, MemoryRateLimiter>> limiterMap = new HashMap<>();

    @Override
    public RateLimiter getRateLimiter(String endpoint, String target, RateLimitRule rule) {
        target = target == null ? TARGET_ALL_KEY : target;
        Map<String, MemoryRateLimiter> map = limiterMap.get(endpoint);
        if (map == null) {
            synchronized (limiterMap) {
                map = limiterMap.computeIfAbsent(endpoint, k -> new HashMap<>());
            }
        }
        MemoryRateLimiter limiter = map.get(target);
        if (limiter == null) {
            synchronized (map) {
                limiter = map.computeIfAbsent(target,
                        k -> new MemoryRateLimiter(rule.getCount(), rule.getPeriod().toMillis()));
            }
        }
        return limiter;
    }
}
