package uk.ac.ebi.phenotype.web.util;

import io.github.bucket4j.*;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.time.Duration;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class PerClientRateLimitInterceptor implements HandlerInterceptor {

    private final Logger log = LoggerFactory.getLogger(this.getClass().getCanonicalName());

    // Cache the most recent 50 hosts and apply rate limiting
    public static final int MAX_SIZE = 50;

    // MAX number of gene page requests from one host in a minute
    public static final long MAX_GENE_PAGE_REQUESTS = 10;

    // MIN number of seconds between gene page requests from one host
    public static final long DELAY_GENE_PAGE_REQUESTS = 2;



    private final Map<String, Bucket> buckets =  Collections.synchronizedMap(new LinkedHashMap<String, Bucket>(MAX_SIZE+1, .75F, true) {
        // This method is called just after a new entry has been added
        public boolean removeEldestEntry(Map.Entry eldest) {
            return size() > MAX_SIZE;
        }
    });


    /**
     * Rate limit the requests to one every 3 seconds and no more than 1 per second
     */
    private static Bucket rateLimitBucket() {
        return Bucket4j.builder()
                .addLimit(
                        Bandwidth.classic(MAX_GENE_PAGE_REQUESTS,
                            Refill.greedy(MAX_GENE_PAGE_REQUESTS, Duration.ofMinutes(1)))
                        .withInitialTokens(MAX_GENE_PAGE_REQUESTS))
                .addLimit(Bandwidth.classic(1,
                            Refill.greedy(1, Duration.ofSeconds(DELAY_GENE_PAGE_REQUESTS))))
                .build();
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response,
                             Object handler) {

        log.debug("Buckets: " + String.join(", ", this.buckets.keySet()));
        String xForwarded = request.getHeader("x-forwarded-host");
        String host = StringUtils.isNotEmpty(xForwarded) ? xForwarded : request.getHeader("host");
        Bucket requestBucket = this.buckets.computeIfAbsent(host, key -> rateLimitBucket());

        log.debug("Using host: " + host + " With bucket: " + requestBucket);

        ConsumptionProbe probe = requestBucket.tryConsumeAndReturnRemaining(1);
        if (probe.isConsumed()) {
            response.addHeader("X-Rate-Limit-Remaining",
                    Long.toString(probe.getRemainingTokens()));
            return true;
        }

        final String waitTime = Long.toString(TimeUnit.NANOSECONDS.toMillis(probe.getNanosToWaitForRefill()));
        log.info("Rate limiting request for page: " + request.getRequestURI() + ", From host: " + host + ", For: " + waitTime + " ms" );

        response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value()); // 429
        response.addHeader("X-Rate-Limit-Retry-After-Milliseconds", waitTime);

        return false;
    }


}
@@ -1041,17 +1041,16 @@
  = 2;%0A%0A%0A
-%0A
     priv
@@ -1088,17 +1088,16 @@
 uckets =
- 
  Collect
@@ -1159,17 +1159,19 @@
 MAX_SIZE
-+
+ + 
 1, .75F,
@@ -1662,32 +1662,36 @@
                 
+    
 Refill.greedy(MA
@@ -1732,24 +1732,32 @@
 inutes(1)))%0A
+        
             
@@ -1859,20 +1859,16 @@
 ssic(1,%0A
-    
         
@@ -2126,24 +2126,165 @@
 handler) %7B%0A%0A
+        if (log.isDebugEnabled()) %7B%0A            // Avoid concurrent modification exception%0A            synchronized (this.buckets) %7B%0A        
         log.
@@ -2342,24 +2342,49 @@
 keySet()));%0A
+            %7D%0A        %7D%0A%0A
         Stri
@@ -2430,20 +2430,19 @@
 rwarded-
-host
+for
 %22);%0A    
@@ -3107,16 +3107,30 @@
 og.info(
+String.format(
 %22Rate li
@@ -3158,92 +3158,150 @@
 ge: 
-%22 + request.getRequestURI() + %22, From host: %22 + host + %22, For: %22 + waitTime + %22 ms%22 
+%25s%5Cn  From host: %25s%5Cn  User-Agent: %25s%5Cn  For: %25sms%22,%0A                request.getRequestURI(), host, request.getHeader(%22User-Agent%22), waitTime)
 );%0A%0A

Fix ConcurrentModificationException while iterating keyset