import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.InputStream;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Map;
import java.util.Objects;

public class Upload {
    static final String ORIGIN = System.getenv().getOrDefault("STASHBAY_ORIGIN", "https://stashbay.net").replaceAll("/$", "");
    static final String KEY = Objects.requireNonNull(System.getenv("STASHBAY_API_KEY"), "Set STASHBAY_API_KEY");
    static final ObjectMapper JSON = new ObjectMapper();
    static final HttpClient CLIENT = HttpClient.newBuilder().followRedirects(HttpClient.Redirect.NEVER).build();

    static JsonNode api(String route, Object body) throws Exception {
        var request = HttpRequest.newBuilder(URI.create(ORIGIN + "/api/v1" + route))
            .timeout(Duration.ofMinutes(5))
            .header("Authorization", "Bearer " + KEY)
            .header("Content-Type", "application/json")
            .POST(HttpRequest.BodyPublishers.ofString(JSON.writeValueAsString(body))).build();
        var response = CLIENT.send(request, HttpResponse.BodyHandlers.ofString());
        check(response);
        return JSON.readTree(response.body());
    }

    static void check(HttpResponse<String> response) {
        if (response.statusCode() < 200 || response.statusCode() >= 300)
            throw new IllegalStateException("HTTP " + response.statusCode() + ": " + response.body());
    }

    public static void main(String[] args) throws Exception {
        if (args.length == 0) throw new IllegalArgumentException("Pass the file path as an argument");
        var path = Path.of(args[0]);
        long size = Files.size(path);
        try (InputStream file = Files.newInputStream(path)) {
            var plan = api("/uploads/init", Map.of("filename", path.getFileName().toString(), "sizeBytes", size));
            var fileId = plan.get("fileId").asText();
            System.err.println("Upload ID: " + fileId);
            int partSize = plan.get("partSize").asInt();
            if (partSize < 1 || partSize > 64 * 1024 * 1024) throw new IllegalStateException("Unexpected part size");
            boolean single = plan.get("single").asBoolean();
            String strategy = plan.get("strategy").asText();
            var parts = new ArrayList<Map<String, Object>>();
            for (int number = 1; number <= plan.get("partCount").asInt(); number++) {
                int length = (int) Math.min(partSize, size - (long) (number - 1) * partSize);
                byte[] bytes = file.readNBytes(length);
                if (bytes.length != length) throw new IllegalStateException("File changed during upload");
                String address = single ? plan.get("uploadUrl").asText() :
                    api("/uploads/parts", Map.of("fileId", fileId, "from", number, "count", 1))
                        .get("parts").get(0).get("url").asText();
                URI base = URI.create(ORIGIN), url = base.resolve(address);
                var request = HttpRequest.newBuilder(url).timeout(Duration.ofMinutes(5));
                if (strategy.equals("proxy") && Objects.equals(url.getScheme(), base.getScheme())
                    && Objects.equals(url.getRawAuthority(), base.getRawAuthority()))
                    request.header("Authorization", "Bearer " + KEY);
                else if (!strategy.equals("presigned") || !url.getScheme().equals("https"))
                    throw new IllegalStateException("Unexpected upload URL or strategy");
                if (single && strategy.equals("presigned")) request.header("If-None-Match", "*");
                var response = CLIENT.send(request.PUT(HttpRequest.BodyPublishers.ofByteArray(bytes)).build(),
                    HttpResponse.BodyHandlers.ofString());
                check(response);
                if (!single) {
                    String etag = strategy.equals("proxy") ? JSON.readTree(response.body()).path("etag").asText("") :
                        response.headers().firstValue("ETag").orElse("");
                    if (etag.isEmpty()) throw new IllegalStateException("Missing part ETag");
                    parts.add(Map.of("partNumber", number, "etag", etag.replace("\"", "")));
                }
            }
            var result = api("/uploads/complete", Map.of("fileId", fileId, "parts", parts));
            System.out.println(result.path("shareUrl").isTextual() ? result.get("shareUrl").asText() : result.toString());
        }
    }
}