import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * Four deterministic, local checks for concurrency mechanisms discussed in the
 * Leaf architecture article. This is teaching code: it does not connect to
 * MySQL or ZooKeeper and does not start Leaf Server.
 *
 * <p>Run with Java 21:</p>
 *
 * <pre>
 * javac LeafMechanismChecks.java
 * java LeafMechanismChecks
 * </pre>
 */
public final class LeafMechanismChecks {
    private static final long TEST_TIMEOUT_SECONDS = 5;

    private LeafMechanismChecks() {
    }

    public static void main(String[] args) throws Exception {
        checkUnsafeCandidateCanMeetReusedUpperBound();
        checkReadLockProtectsCandidateAndUpperBound();
        checkCallerRunsCannotUpgradeReadLock();
        checkRejectedSubmissionRestoresLoadingFlag();
        System.out.println("All 4 Leaf mechanism checks passed.");
    }

    /**
     * Without request-side locking, a candidate from an exhausted generation
     * can pause while the same mutable segment object is reused. The stale
     * candidate is then compared with the new generation's upper bound.
     */
    private static void checkUnsafeCandidateCanMeetReusedUpperBound() throws Exception {
        ReusableSegment segment = new ReusableSegment(200, 200);
        CountDownLatch candidateRead = new CountDownLatch(1);
        CountDownLatch segmentReused = new CountDownLatch(1);
        ExecutorService executor = Executors.newFixedThreadPool(2);

        try {
            Future<Observation> request = executor.submit(() -> {
                long candidate = segment.value.getAndIncrement();
                candidateRead.countDown();
                await(segmentReused, "segment reuse");
                long observedUpperBound = segment.upperExclusive;
                return new Observation(
                        candidate,
                        observedUpperBound,
                        candidate < observedUpperBound);
            });

            Future<Void> recycler = executor.submit(() -> {
                await(candidateRead, "unsafe candidate read");
                // Simulate two buffer switches and reuse of the old object for
                // the later [400, 500) generation.
                segment.value.set(400);
                segment.upperExclusive = 500;
                segmentReused.countDown();
                return null;
            });

            Observation observation = request.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
            recycler.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);

            require(observation.candidate == 200, "expected stale candidate 200");
            require(observation.observedUpperBound == 500, "expected reused upper bound 500");
            require(observation.accepted, "unsafe model should incorrectly accept candidate 200");

            System.out.printf(
                    "PASS 1/4 unsafe mix: candidate=%d, observedUpperBound=%d, accepted=%s%n",
                    observation.candidate,
                    observation.observedUpperBound,
                    observation.accepted);
        } finally {
            segmentReused.countDown();
            shutdownNow(executor);
        }
    }

    /**
     * With a read lock around candidate acquisition and bound checking, a
     * writer cannot switch/reuse the segment until the exhausted candidate has
     * been rejected against its original bound.
     */
    private static void checkReadLockProtectsCandidateAndUpperBound() throws Exception {
        ReusableSegment segment = new ReusableSegment(200, 200);
        ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
        CountDownLatch readerHasCandidate = new CountDownLatch(1);
        CountDownLatch writerIsAttempting = new CountDownLatch(1);
        CountDownLatch writerAcquired = new CountDownLatch(1);
        CountDownLatch allowReaderToCompare = new CountDownLatch(1);
        ExecutorService executor = Executors.newFixedThreadPool(2);

        try {
            Future<Observation> request = executor.submit(() -> {
                lock.readLock().lock();
                try {
                    long candidate = segment.value.getAndIncrement();
                    readerHasCandidate.countDown();
                    await(allowReaderToCompare, "reader comparison release");
                    long observedUpperBound = segment.upperExclusive;
                    return new Observation(
                            candidate,
                            observedUpperBound,
                            candidate < observedUpperBound);
                } finally {
                    lock.readLock().unlock();
                }
            });

            Future<Void> writer = executor.submit(() -> {
                await(readerHasCandidate, "protected candidate read");
                writerIsAttempting.countDown();
                lock.writeLock().lock();
                try {
                    writerAcquired.countDown();
                    segment.value.set(400);
                    segment.upperExclusive = 500;
                } finally {
                    lock.writeLock().unlock();
                }
                return null;
            });

            await(readerHasCandidate, "reader acquisition");
            await(writerIsAttempting, "writer attempt");
            boolean writerFinishedProtectedSwitch = writerAcquired.await(200, TimeUnit.MILLISECONDS);
            require(!writerFinishedProtectedSwitch,
                    "writer must not acquire the write lock while the reader holds its lock");

            allowReaderToCompare.countDown();
            Observation observation = request.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
            writer.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);

            require(observation.candidate == 200, "expected exhausted candidate 200");
            require(observation.observedUpperBound == 200, "reader must observe original upper bound 200");
            require(!observation.accepted, "protected exhausted candidate must be rejected");
            require(writerAcquired.getCount() == 0, "writer should proceed after the read lock is released");

            System.out.printf(
                    "PASS 2/4 read lock: writerBlocked=%s, candidate=%d, observedUpperBound=%d, accepted=%s%n",
                    !writerFinishedProtectedSwitch,
                    observation.candidate,
                    observation.observedUpperBound,
                    observation.accepted);
        } finally {
            allowReaderToCompare.countDown();
            shutdownNow(executor);
        }
    }

    /**
     * Saturate a bounded executor so CallerRunsPolicy executes the refill task
     * on the request thread. Because that thread still holds the read lock, a
     * timed attempt to obtain the write lock cannot perform an unsupported
     * read-to-write upgrade.
     */
    private static void checkCallerRunsCannotUpgradeReadLock() throws Exception {
        ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
        CountDownLatch workerStarted = new CountDownLatch(1);
        CountDownLatch releaseWorker = new CountDownLatch(1);
        CountDownLatch queuedTaskFinished = new CountDownLatch(1);
        ThreadPoolExecutor executor = new ThreadPoolExecutor(
                1,
                1,
                0L,
                TimeUnit.MILLISECONDS,
                new ArrayBlockingQueue<>(1),
                new ThreadPoolExecutor.CallerRunsPolicy());

        AtomicBoolean ranOnRequestThread = new AtomicBoolean(false);
        AtomicBoolean writeLockAcquired = new AtomicBoolean(true);
        Thread requestThread = Thread.currentThread();

        try {
            executor.execute(() -> {
                workerStarted.countDown();
                awaitUninterruptibly(releaseWorker);
            });
            await(workerStarted, "executor worker start");

            // The worker is occupied, so this task fills the sole queue slot.
            executor.execute(queuedTaskFinished::countDown);

            lock.readLock().lock();
            try {
                // Pool and queue are both full. CallerRunsPolicy executes this
                // refill task synchronously on requestThread.
                executor.execute(() -> {
                    ranOnRequestThread.set(Thread.currentThread() == requestThread);
                    boolean acquired = false;
                    try {
                        acquired = lock.writeLock().tryLock(200, TimeUnit.MILLISECONDS);
                        writeLockAcquired.set(acquired);
                    } catch (InterruptedException exception) {
                        Thread.currentThread().interrupt();
                        throw new IllegalStateException("write-lock check interrupted", exception);
                    } finally {
                        if (acquired) {
                            lock.writeLock().unlock();
                        }
                    }
                });
            } finally {
                lock.readLock().unlock();
            }

            require(ranOnRequestThread.get(), "CallerRunsPolicy should use the request thread");
            require(!writeLockAcquired.get(),
                    "a request thread holding the read lock must not acquire the write lock");

            System.out.printf(
                    "PASS 3/4 CallerRunsPolicy: ranOnRequestThread=%s, writeLockAcquired=%s%n",
                    ranOnRequestThread.get(),
                    writeLockAcquired.get());
        } finally {
            releaseWorker.countDown();
            executor.shutdown();
            require(executor.awaitTermination(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS),
                    "CallerRunsPolicy executor did not terminate");
            require(queuedTaskFinished.getCount() == 0, "queued task did not finish");
        }
    }

    /**
     * If loading=true is set before execute(), a rejected submission never
     * reaches the task's finally block. The submitter must therefore restore
     * the flag in the rejection path.
     */
    private static void checkRejectedSubmissionRestoresLoadingFlag() throws Exception {
        AtomicBoolean loading = new AtomicBoolean(false);
        AtomicBoolean taskRan = new AtomicBoolean(false);
        ThreadPoolExecutor executor = new ThreadPoolExecutor(
                1,
                1,
                0L,
                TimeUnit.MILLISECONDS,
                new ArrayBlockingQueue<>(1),
                new ThreadPoolExecutor.AbortPolicy());
        executor.shutdown();

        require(loading.compareAndSet(false, true), "failed to acquire loading flag");
        boolean rejected = false;
        try {
            executor.execute(() -> {
                taskRan.set(true);
                try {
                    // A real task would load a segment here.
                } finally {
                    loading.set(false);
                }
            });
        } catch (RejectedExecutionException expected) {
            rejected = true;
            // The task did not run, so its finally block cannot restore state.
            loading.set(false);
        }

        require(rejected, "submission should be rejected after executor shutdown");
        require(!taskRan.get(), "rejected task must not run");
        require(!loading.get(), "submitter must restore loading=false after rejection");
        require(executor.awaitTermination(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS),
                "rejected-submission executor did not terminate");

        System.out.printf(
                "PASS 4/4 rejected submission: rejected=%s, taskRan=%s, loading=%s%n",
                rejected,
                taskRan.get(),
                loading.get());
    }

    private static void await(CountDownLatch latch, String description) throws InterruptedException {
        if (!latch.await(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
            throw new AssertionError("timed out waiting for " + description);
        }
    }

    private static void awaitUninterruptibly(CountDownLatch latch) {
        boolean interrupted = false;
        while (true) {
            try {
                latch.await();
                break;
            } catch (InterruptedException exception) {
                interrupted = true;
            }
        }
        if (interrupted) {
            Thread.currentThread().interrupt();
        }
    }

    private static void shutdownNow(ExecutorService executor) throws InterruptedException {
        executor.shutdownNow();
        require(executor.awaitTermination(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS),
                "executor did not terminate");
    }

    private static void require(boolean condition, String message) {
        if (!condition) {
            throw new AssertionError(message);
        }
    }

    private record Observation(long candidate, long observedUpperBound, boolean accepted) {
    }

    private static final class ReusableSegment {
        private final AtomicLong value;
        private volatile long upperExclusive;

        private ReusableSegment(long nextValue, long upperExclusive) {
            this.value = new AtomicLong(nextValue);
            this.upperExclusive = upperExclusive;
        }
    }
}