import java.util.concurrent.CountDownLatch;
import java.util.concurrent.StructuredTaskScope;
import java.util.concurrent.atomic.AtomicLong;

/**
 * JEP 506's own motivation for Scoped Values over ThreadLocal names "expensive inheritance" as
 * one of ThreadLocal's three flaws: a child thread must allocate storage for every inheritable
 * thread-local value copied from its parent. A Scoped Value binding isn't copied into the child
 * at all -- as InheritanceRequiresStructuredTaskScope.java shows, a child only sees a binding if
 * it's forked from a StructuredTaskScope, and what it gets is a reference to the SAME immutable
 * binding chain its parent already holds, not a copy.
 *
 * This measures that difference directly instead of repeating the claim: N subtasks are forked
 * (via StructuredTaskScope, since that's the only mechanism that inherits a Scoped Value binding
 * at all -- see the other file), each holding 5 bound values (ScopedValue in one run,
 * InheritableThreadLocal in the other), and parked simultaneously so all N are alive and holding
 * their bindings at once. Heap used is measured immediately before start and immediately after
 * every subtask confirms it has read its values -- the delta is what N threads' worth of "5 bound
 * values, held live" actually costs.
 *
 * Compile and run with --enable-preview --release 25 (StructuredTaskScope is still preview on 25).
 * Chapter: docs/17-scoped-values-migration.md
 */
public class ScopedValueVsThreadLocalMemory {

    static final ScopedValue<String> SV1 = ScopedValue.newInstance();
    static final ScopedValue<String> SV2 = ScopedValue.newInstance();
    static final ScopedValue<String> SV3 = ScopedValue.newInstance();
    static final ScopedValue<String> SV4 = ScopedValue.newInstance();
    static final ScopedValue<String> SV5 = ScopedValue.newInstance();

    static final ThreadLocal<String> TL1 = new InheritableThreadLocal<>();
    static final ThreadLocal<String> TL2 = new InheritableThreadLocal<>();
    static final ThreadLocal<String> TL3 = new InheritableThreadLocal<>();
    static final ThreadLocal<String> TL4 = new InheritableThreadLocal<>();
    static final ThreadLocal<String> TL5 = new InheritableThreadLocal<>();

    static final long EXPECTED_CHECKSUM_PER_THREAD = 5 + 5 + 7 + 5 + 4; // "alpha".."echo" lengths

    public static void main(String[] args) throws Exception {
        int n = args.length > 0 ? Integer.parseInt(args[0]) : 200_000;
        System.out.println("threads per scenario: " + n + ", 5 bound values each, forked via StructuredTaskScope");

        // Warm up classes/JIT with a small run so the real measurement doesn't pay class-loading cost.
        runScopedValueScenario(1_000);
        runThreadLocalScenario(1_000);

        long scopedBytes = runScopedValueScenario(n);
        System.out.println("ScopedValue:  " + n + " live subtasks, 5 bound values each, heap delta = "
                + fmt(scopedBytes) + " (" + (scopedBytes / n) + " bytes/thread)");

        long threadLocalBytes = runThreadLocalScenario(n);
        System.out.println("ThreadLocal:  " + n + " live subtasks, 5 inherited values each, heap delta = "
                + fmt(threadLocalBytes) + " (" + (threadLocalBytes / n) + " bytes/thread)");

        double ratio = (double) threadLocalBytes / scopedBytes;
        System.out.printf("ThreadLocal used %.1fx the heap ScopedValue used, for holding the identical 5 values live on %d threads%n", ratio, n);
        // Deliberately NOT asserted as threadLocalBytes > scopedBytes: total-heap-before/after is
        // too coarse to trust at this scale. The VirtualThread + StackChunk + SubtaskImpl objects
        // that exist either way run 1000-2000+ bytes/thread; the actual difference this file is
        // trying to isolate -- InheritableThreadLocal's per-thread ThreadLocalMap -- only comes to
        // roughly 260 bytes/thread (measured precisely in MemoryFootprintProbe.java's output, via
        // jcmd's exact object census, not GC-timing-dependent Runtime.freeMemory()). A ~15% signal
        // sitting inside a multi-hundred-MB heap, remeasured across three System.gc() calls, is
        // exactly the kind of thing that can and did flip sign between runs here. Leaving this
        // file's honest, noisy number in place rather than deleting it: it's the reason this repo
        // reaches for GC.class_histogram in the next file instead of trusting this one alone.
    }

    @SuppressWarnings("preview")
    static long runScopedValueScenario(int n) throws Exception {
        CountDownLatch ready = new CountDownLatch(n);
        CountDownLatch release = new CountDownLatch(1);
        AtomicLong checksum = new AtomicLong();

        long before = heapUsedAfterGc();
        return ScopedValue.where(SV1, "alpha").where(SV2, "bravo").where(SV3, "charlie").where(SV4, "delta").where(SV5, "echo")
                .call(() -> {
                    try (var scope = StructuredTaskScope.open()) {
                        for (int i = 0; i < n; i++) {
                            scope.fork(() -> {
                                checksum.addAndGet(SV1.get().length() + SV2.get().length() + SV3.get().length() + SV4.get().length() + SV5.get().length());
                                ready.countDown();
                                await(release);
                                return null;
                            });
                        }
                        ready.await();
                        long after = heapUsedAfterGc();
                        release.countDown();
                        scope.join();
                        Check.that(checksum.get() == (long) n * EXPECTED_CHECKSUM_PER_THREAD, "every one of " + n + " subtasks read all 5 scoped values correctly");
                        return after - before;
                    }
                });
    }

    @SuppressWarnings("preview")
    static long runThreadLocalScenario(int n) throws Exception {
        CountDownLatch ready = new CountDownLatch(n);
        CountDownLatch release = new CountDownLatch(1);
        AtomicLong checksum = new AtomicLong();

        TL1.set("alpha");
        TL2.set("bravo");
        TL3.set("charlie");
        TL4.set("delta");
        TL5.set("echo");

        long before = heapUsedAfterGc();
        long delta;
        try (var scope = StructuredTaskScope.open()) {
            for (int i = 0; i < n; i++) {
                scope.fork(() -> {
                    checksum.addAndGet(TL1.get().length() + TL2.get().length() + TL3.get().length() + TL4.get().length() + TL5.get().length());
                    ready.countDown();
                    await(release);
                    return null;
                });
            }
            ready.await();
            long after = heapUsedAfterGc();
            delta = after - before;
            release.countDown();
            scope.join();
        }
        TL1.remove(); TL2.remove(); TL3.remove(); TL4.remove(); TL5.remove();
        Check.that(checksum.get() == (long) n * EXPECTED_CHECKSUM_PER_THREAD, "every one of " + n + " subtasks read all 5 inherited ThreadLocal values correctly");
        return delta;
    }

    static void await(CountDownLatch latch) {
        try {
            latch.await();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    static long heapUsedAfterGc() throws InterruptedException {
        for (int i = 0; i < 3; i++) {
            System.gc();
            Thread.sleep(200);
        }
        Runtime rt = Runtime.getRuntime();
        return rt.totalMemory() - rt.freeMemory();
    }

    static String fmt(long bytes) {
        return String.format("%.1f MB", bytes / (1024.0 * 1024.0));
    }
}