From cdd86fe8ceb9c471584eb562f8d0ae63753f8d35 Mon Sep 17 00:00:00 2001
From: JoshiCodes <55353244+JoshiCodes@users.noreply.github.com>
Date: Tue, 13 Aug 2024 17:27:59 +0200
Subject: [PATCH] ImageAction#execute and RestAction#execute now waits after a
 rate limit (Mentioned in #2)

---
 .../aghpb4j/action/ImageAction.java           | 11 ++--
 .../joshicodes/aghpb4j/action/RestAction.java | 53 +++++++++++++++----
 2 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/src/main/java/de/joshicodes/aghpb4j/action/ImageAction.java b/src/main/java/de/joshicodes/aghpb4j/action/ImageAction.java
index c18baef..0fb62ac 100644
--- a/src/main/java/de/joshicodes/aghpb4j/action/ImageAction.java
+++ b/src/main/java/de/joshicodes/aghpb4j/action/ImageAction.java
@@ -1,6 +1,5 @@
 package de.joshicodes.aghpb4j.action;
 
-import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
@@ -18,16 +17,12 @@ public ImageAction(String url, Class<T> tClass, Function<HttpRequest.Builder, Ht
     @Override
     public T execute() {
         try {
-            HttpClient.Builder client = HttpClient.newBuilder()
-                    .followRedirects(HttpClient.Redirect.NORMAL)
-                    .version(HttpClient.Version.HTTP_2);
-            HttpRequest.Builder request = HttpRequest.newBuilder()
-                    .uri(URI.create(url))
-                    .method(method, HttpRequest.BodyPublishers.noBody());
+            HttpClient.Builder client = build();
+            HttpRequest.Builder request = buildRequest();
             if(clientModifier != null) {
                 request = clientModifier.apply(request);
             }
-            HttpResponse<byte[]> response = client.build().send(request.build(), HttpResponse.BodyHandlers.ofByteArray());
+            HttpResponse<byte[]> response = sendRequest(client.build(), request.build(), HttpResponse.BodyHandlers.ofByteArray(), 3);
             return this.responseHandler.apply(new RestResponse<>(response, byte[].class));
         } catch (Exception e) {
             throw new RuntimeException(e);
diff --git a/src/main/java/de/joshicodes/aghpb4j/action/RestAction.java b/src/main/java/de/joshicodes/aghpb4j/action/RestAction.java
index 32fb786..fd4eac5 100644
--- a/src/main/java/de/joshicodes/aghpb4j/action/RestAction.java
+++ b/src/main/java/de/joshicodes/aghpb4j/action/RestAction.java
@@ -96,22 +96,55 @@ public void queue(Consumer<T> success, Consumer<Throwable> failure) {
      */
     public T execute() {
         try {
-            HttpClient.Builder client = HttpClient.newBuilder()
-                    .followRedirects(HttpClient.Redirect.NORMAL)
-                    .version(HttpClient.Version.HTTP_2);
-            HttpRequest.Builder request = HttpRequest.newBuilder()
-                    .uri(URI.create(url))
-                    .method(method, HttpRequest.BodyPublishers.noBody());
-            if(clientModifier != null) {
-                request = clientModifier.apply(request);
-            }
-            HttpResponse<String> response = client.build().send(request.build(), HttpResponse.BodyHandlers.ofString());
+            HttpClient.Builder client = build();
+            HttpRequest.Builder request = buildRequest();
+            HttpResponse<String> response = sendRequest(client.build(), request.build(), HttpResponse.BodyHandlers.ofString(), 3);
             return responseHandler.apply(new RestResponse<>(response, String.class));
         } catch (Exception e) {
             throw new RuntimeException(e);
         }
     }
 
+    protected <C> HttpResponse<C> sendRequest(HttpClient client, HttpRequest request, HttpResponse.BodyHandler<C> handler, int retries) {
+        try {
+            HttpResponse<C> response = client.send(request, handler);
+            if(response.headers().firstValue("x-ratelimit-remaining").orElse("1").equals("0")) {
+                // Rate limit
+                final long reset = Long.parseLong(response.headers().firstValue("x-ratelimit-reset").orElse("0"));
+                final long now = System.currentTimeMillis() / 1000;
+                final long diff = reset - now;
+                if(diff > 0) {
+                    // Sleep
+                    Thread.sleep(diff * 1000);
+                    // Retry
+                    return sendRequest(client, request, handler, retries - 1);
+                }
+            }
+            return response;
+        } catch (Exception e) {
+            if(retries > 0) {
+                return sendRequest(client, request, handler, retries - 1);
+            }
+            throw new RuntimeException(e);
+        }
+    }
+
+    protected HttpClient.Builder build() {
+        return HttpClient.newBuilder()
+                .followRedirects(HttpClient.Redirect.NORMAL)
+                .version(HttpClient.Version.HTTP_2);
+    }
+
+    protected HttpRequest.Builder buildRequest() {
+        HttpRequest.Builder request = HttpRequest.newBuilder()
+                .uri(URI.create(url))
+                .method(method, HttpRequest.BodyPublishers.noBody());
+        if(clientModifier != null) {
+            request = clientModifier.apply(request);
+        }
+        return request;
+    }
+
     public record RestResponse<A>(HttpResponse<A> httpResponse, Class<A> aClass) {
 
         public String rawBody() {