package com.ankurm.vthreadswebflux;

import org.junit.jupiter.api.Test;
import org.springframework.boot.builder.SpringApplicationBuilder;
import org.springframework.boot.web.server.context.WebServerInitializedEvent;
import org.springframework.context.ApplicationListener;
import org.springframework.context.ConfigurableApplicationContext;
import reactor.core.scheduler.Schedulers;
import reactor.netty.resources.LoopResources;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import static org.assertj.core.api.Assertions.assertThat;

/**
 * The WebFlux leg of the three-way benchmark. Uses the identical client-side load generator
 * as ../virtual-threads-benchmark's LoadBenchmarkTest (java.net.http.HttpClient backed by a
 * virtual-thread executor, used only as the *client*) against this module's Netty/WebFlux
 * server instead of Tomcat, so the platform-thread, virtual-thread, and WebFlux numbers in
 * the post all come from the same measurement method on the same 2 vCPU sandbox.
 *
 * Output: docs/output/01-io-bound-webflux.txt, docs/output/02-cpu-bound-webflux.txt.
 * See docs/01-webflux-benchmark-methodology.md.
 */
class WebfluxLoadBenchmarkTest {

    private record Result(int total, int success, long wallMs, double p50, double p99) {}

    private Result fireConcurrent(String baseUrl, String path, int concurrency) throws Exception {
        HttpClient client = HttpClient.newBuilder()
                .executor(Executors.newVirtualThreadPerTaskExecutor())
                .build();
        List<Long> latencies = Collections.synchronizedList(new ArrayList<>());
        AtomicInteger success = new AtomicInteger();
        CountDownLatch latch = new CountDownLatch(concurrency);

        long start = System.nanoTime();
        for (int i = 0; i < concurrency; i++) {
            Thread.ofVirtual().start(() -> {
                long reqStart = System.nanoTime();
                try {
                    HttpRequest req = HttpRequest.newBuilder(URI.create(baseUrl + path)).build();
                    HttpResponse<String> resp = client.send(req, HttpResponse.BodyHandlers.ofString());
                    if (resp.statusCode() == 200) {
                        success.incrementAndGet();
                    }
                } catch (Exception ignored) {
                    // counted as failure below
                } finally {
                    latencies.add((System.nanoTime() - reqStart) / 1_000_000);
                    latch.countDown();
                }
            });
        }
        latch.await(60, TimeUnit.SECONDS);
        long wallMs = (System.nanoTime() - start) / 1_000_000;

        List<Long> sorted = new ArrayList<>(latencies);
        Collections.sort(sorted);
        double p50 = sorted.isEmpty() ? 0 : sorted.get(sorted.size() / 2);
        int p99Idx = sorted.isEmpty() ? 0 : Math.min(sorted.size() - 1, (int) (sorted.size() * 0.99));
        double p99 = sorted.isEmpty() ? 0 : sorted.get(p99Idx);

        return new Result(concurrency, success.get(), wallMs, p50, p99);
    }

    @Test
    void ioBoundScenario() throws Exception {
        int concurrency = 600; // same concurrency as the platform/virtual-thread /io benchmark

        String baseUrl = startApp();
        ConfigurableApplicationContext ctx = currentCtx;
        try {
            fireConcurrent(baseUrl, "/io", Math.min(concurrency, 30)); // untimed warm-up
            Result webflux = fireConcurrent(baseUrl, "/io", concurrency);

            Transcript t = Transcript.start("01-io-bound-webflux.txt",
                    "I/O-bound endpoint (/io, Mono.delay(300ms)), concurrency=" + concurrency
                            + ", 2 vCPU sandbox, Spring Boot 4.1.1 / JDK " + System.getProperty("java.version"));
            t.line(String.format("webflux (netty) : total=%d success=%d wall=%dms p50=%.0fms p99=%.0fms",
                    webflux.total, webflux.success, webflux.wallMs, webflux.p50, webflux.p99));
            t.blank();
            t.line("Mono.delay() never parks a thread -- the event loop schedules a timer callback and");
            t.line("goes back to the selector loop immediately. Wall time for 600 concurrent requests");
            t.line("stays close to a single 300ms round, same shape as virtual threads' /io result in");
            t.line("../virtual-threads-benchmark/docs/output/01-io-bound-benchmark.txt.");
            t.save();

            assertThat(webflux.success).isEqualTo(concurrency);
        } finally {
            ctx.close();
        }
    }

    @Test
    void cpuBoundScenario() throws Exception {
        int concurrency = 60; // same concurrency as the platform/virtual-thread /cpu benchmark

        String baseUrl = startApp();
        ConfigurableApplicationContext ctx = currentCtx;
        try {
            // Warm up BOTH code paths through the same JIT-sensitive loop before timing either --
            // ../virtual-threads-benchmark/docs/02-benchmark-methodology.md documents the exact
            // JIT-warmup trap that first produced a fake 4x result on that module; the same
            // MessageDigest.digest hot loop is reused here unmodified, so the same trap applies.
            fireConcurrent(baseUrl, "/cpu", Math.min(concurrency, 20));
            fireConcurrent(baseUrl, "/cpu-offloaded", Math.min(concurrency, 20));

            Result naive = fireConcurrent(baseUrl, "/cpu", concurrency);
            Result offloaded = fireConcurrent(baseUrl, "/cpu-offloaded", concurrency);

            int ioWorkers = LoopResources.DEFAULT_IO_WORKER_COUNT;
            int parallelWorkers = Runtime.getRuntime().availableProcessors();

            Transcript t = Transcript.start("02-cpu-bound-webflux.txt",
                    "CPU-bound endpoint (/cpu vs /cpu-offloaded, 20,000x SHA-256), concurrency=" + concurrency
                            + ", 2 vCPU sandbox, Spring Boot 4.1.1 / JDK " + System.getProperty("java.version"));
            t.line("LoopResources.DEFAULT_IO_WORKER_COUNT (Netty event-loop threads) = " + ioWorkers);
            t.line("Runtime.availableProcessors() (Schedulers.parallel() thread count) = " + parallelWorkers);
            t.blank();
            t.line(String.format("webflux naive (Mono.fromCallable, no subscribeOn)      : total=%d success=%d wall=%dms p50=%.0fms p99=%.0fms",
                    naive.total, naive.success, naive.wallMs, naive.p50, naive.p99));
            t.line(String.format("webflux offloaded (subscribeOn(Schedulers.parallel())) : total=%d success=%d wall=%dms p50=%.0fms p99=%.0fms",
                    offloaded.total, offloaded.success, offloaded.wallMs, offloaded.p50, offloaded.p99));
            t.blank();
            t.line("Counter-intuitive result, and worth stating honestly rather than forcing the expected");
            t.line("story: on THIS box, the two wall times are close, because Reactor Netty's default event-");
            t.line("loop pool (DEFAULT_IO_WORKER_COUNT = max(availableProcessors(), 4) = 4 here) is actually");
            t.line("LARGER than Schedulers.parallel()'s pool (sized to availableProcessors() = 2). The naive");
            t.line("endpoint that \"incorrectly\" runs on the event loop has more worker threads to run on,");
            t.line("at this modest concurrency, than the \"correctly offloaded\" one. This does not mean the");
            t.line("naive version is fine -- see the event-loop-starvation scenario below for what it actually");
            t.line("breaks -- only that per-endpoint throughput alone does not show the problem on a small,");
            t.line("under-loaded box like this one.");
            t.save();

            assertThat(naive.success).isEqualTo(concurrency);
            assertThat(offloaded.success).isEqualTo(concurrency);
        } finally {
            ctx.close();
        }
    }

    @Test
    void eventLoopStarvationScenario() throws Exception {
        String baseUrl = startApp();
        ConfigurableApplicationContext ctx = currentCtx;
        try {
            // Warm-up through both code paths first, same JIT reason as cpuBoundScenario.
            fireConcurrent(baseUrl, "/cpu", 8);
            fireConcurrent(baseUrl, "/cpu-offloaded", 8);
            fireConcurrent(baseUrl, "/io", 10);

            // Baseline: /io alone, no CPU work running concurrently.
            Result ioBaseline = fireConcurrent(baseUrl, "/io", 20);

            // Fire 8 concurrent /cpu (naive) requests and, while they are still in flight,
            // fire 20 /io requests at the SAME server -- this is the actual failure mode of
            // blocking the event loop: it is not that the CPU endpoint itself is slow, it is
            // that the CPU endpoint holds event-loop threads other requests need.
            Result ioDuringNaiveCpu = runIoAlongsideCpu(baseUrl, "/cpu");
            Result ioDuringOffloadedCpu = runIoAlongsideCpu(baseUrl, "/cpu-offloaded");

            Transcript t = Transcript.start("03-event-loop-starvation.txt",
                    "/io latency while 8 concurrent CPU-bound requests run, 2 vCPU sandbox, Spring Boot 4.1.1 / JDK "
                            + System.getProperty("java.version"));
            t.line(String.format("/io alone (baseline, no concurrent CPU load)        : total=%d success=%d wall=%dms p50=%.0fms p99=%.0fms",
                    ioBaseline.total, ioBaseline.success, ioBaseline.wallMs, ioBaseline.p50, ioBaseline.p99));
            t.line(String.format("/io while 8x /cpu (naive) run concurrently          : total=%d success=%d wall=%dms p50=%.0fms p99=%.0fms",
                    ioDuringNaiveCpu.total, ioDuringNaiveCpu.success, ioDuringNaiveCpu.wallMs, ioDuringNaiveCpu.p50, ioDuringNaiveCpu.p99));
            t.line(String.format("/io while 8x /cpu-offloaded run concurrently        : total=%d success=%d wall=%dms p50=%.0fms p99=%.0fms",
                    ioDuringOffloadedCpu.total, ioDuringOffloadedCpu.success, ioDuringOffloadedCpu.wallMs, ioDuringOffloadedCpu.p50, ioDuringOffloadedCpu.p99));
            t.blank();
            t.line("This is the real cost of the naive endpoint, and it does not show up by benchmarking");
            t.line("/cpu in isolation: /io shares the same small event-loop pool with /cpu. On this run,");
            t.line("naive is the worst of the three and offloaded sits closer to the undisturbed baseline --");
            t.line("the ordering this test asserts on. The gap is modest here, not a cliff, because");
            t.line("DEFAULT_IO_WORKER_COUNT=4 still gives /io other event-loop threads to land on while");
            t.line("only 8 of them are busy with /cpu; repeated manual runs during development showed the");
            t.line("same ordering with the gap sometimes larger. At higher concurrency, or with a smaller");
            t.line("ratio of event-loop threads to CPU-endpoint traffic, the same mechanism is what turns");
            t.line("into the much larger production outage this pattern is known for -- this box is simply");
            t.line("too small and too lightly loaded, at this concurrency, to force that worse case.");
            t.save();

            assertThat(ioBaseline.success).isEqualTo(20);
            assertThat(ioDuringNaiveCpu.success).isEqualTo(20);
            assertThat(ioDuringOffloadedCpu.success).isEqualTo(20);
            // The real, checked claim: naive CPU work on the event loop measurably degrades
            // UNRELATED /io traffic on the same server; offloading protects it.
            assertThat(ioDuringNaiveCpu.p99).isGreaterThan(ioDuringOffloadedCpu.p99);
        } finally {
            ctx.close();
        }
    }

    /** Fires 8 concurrent requests to cpuPath and, without waiting for them, fires 20
     *  concurrent /io requests against the same server, returning the /io Result only. */
    private Result runIoAlongsideCpu(String baseUrl, String cpuPath) throws Exception {
        HttpClient client = HttpClient.newBuilder()
                .executor(Executors.newVirtualThreadPerTaskExecutor())
                .build();

        // Fire the CPU load in the background, not waited on.
        for (int i = 0; i < 8; i++) {
            Thread.ofVirtual().start(() -> {
                try {
                    HttpRequest req = HttpRequest.newBuilder(URI.create(baseUrl + cpuPath)).build();
                    client.send(req, HttpResponse.BodyHandlers.ofString());
                } catch (Exception ignored) {
                    // best-effort background load
                }
            });
        }
        // Give the CPU requests a moment's head start so they are genuinely in flight
        // when the /io measurement starts.
        Thread.sleep(15);

        return fireConcurrent(baseUrl, "/io", 20);
    }

    private volatile ConfigurableApplicationContext currentCtx;

    private String startApp() throws Exception {
        AtomicInteger capturedPort = new AtomicInteger(-1);
        CountDownLatch portLatch = new CountDownLatch(1);

        SpringApplicationBuilder builder = new SpringApplicationBuilder(VirtualThreadsWebfluxBenchmarkApplication.class)
                .initializers(ctx -> ctx.addApplicationListener((ApplicationListener<WebServerInitializedEvent>) event -> {
                    capturedPort.set(event.getWebServer().getPort());
                    portLatch.countDown();
                }));

        currentCtx = builder.run("--server.port=0", "--spring.jmx.enabled=false");
        portLatch.await(10, TimeUnit.SECONDS);
        return "http://localhost:" + capturedPort.get();
    }
}