/*
 * Copyright (c) 2017, 2018, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

import java.nio.ByteBuffer;
import java.util.List;
import java.util.Random;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Flow;
import java.util.concurrent.Flow.Subscription;
import java.util.concurrent.SubmissionPublisher;
import java.util.function.BiConsumer;
import java.net.http.HttpResponse.BodyHandler;
import java.net.http.HttpResponse.BodyHandlers;
import java.net.http.HttpResponse.BodySubscriber;
import java.net.http.HttpResponse.BodySubscribers;
import jdk.test.lib.RandomFactory;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static java.lang.Long.MAX_VALUE;
import static java.lang.Long.min;
import static java.lang.System.out;
import static java.util.concurrent.CompletableFuture.delayedExecutor;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static org.testng.Assert.*;

/*
 * @test
 * @bug 8184285
 * @summary Direct test for HttpResponse.BodySubscriber.buffering() API
 * @key randomness
 * @library /test/lib
 * @build jdk.test.lib.RandomFactory
 * @run testng/othervm -Djdk.internal.httpclient.debug=true BufferingSubscriberTest
 */

public class BufferingSubscriberTest {

    // If we compute that a test will take less that 10s
    // we judge it acceptable
    static final long LOWER_THRESHOLD = 10_000; // 10 sec.
    // If we compute that a test will take more than 20 sec
    // we judge it problematic: we will try to adjust the
    // buffer sizes, and if we can't we will print a warning
    static final long UPPER_THRESHOLD = 20_000; // 20 sec.

    static final Random random = RandomFactory.getRandom();
    static final long start = System.nanoTime();
    static final String START = "start";
    static final String END   = "end  ";
    static long elapsed() { return (System.nanoTime() - start)/1000_000;}
    static void printStamp(String what, String fmt, Object... args) {
        long elapsed = elapsed();
        long sec = elapsed/1000;
        long ms  = elapsed % 1000;
        String time = sec > 0 ? sec + "sec " : "";
        time = time + ms + "ms";
        out.println(what + "\t ["+time+"]\t "+ String.format(fmt,args));
    }
    @DataProvider(name = "negatives")
    public Object[][] negatives() {
        return new Object[][] {  { 0 }, { -1 }, { -1000 } };
    }

    @Test(dataProvider = "negatives", expectedExceptions = IllegalArgumentException.class)
    public void subscriberThrowsIAE(int bufferSize) {
        printStamp(START, "subscriberThrowsIAE(%d)", bufferSize);
        try {
            BodySubscriber<?> bp = BodySubscribers.ofByteArray();
            BodySubscribers.buffering(bp, bufferSize);
        } finally {
            printStamp(END, "subscriberThrowsIAE(%d)", bufferSize);
        }
    }

    @Test(dataProvider = "negatives", expectedExceptions = IllegalArgumentException.class)
    public void handlerThrowsIAE(int bufferSize) {
        printStamp(START, "handlerThrowsIAE(%d)", bufferSize);
        try {
            BodyHandler<?> bp = BodyHandlers.ofByteArray();
            BodyHandlers.buffering(bp, bufferSize);
        } finally {
            printStamp(END, "handlerThrowsIAE(%d)", bufferSize);
        }
    }

    // ---

    @DataProvider(name = "config")
    public Object[][] config() {
        return new Object[][] {
            // iterations delayMillis numBuffers bufferSize maxBufferSize minBufferSize
            { 1,              0,          1,         1,         2,            1   },
            { 1,              5,          1,         100,       2,            1   },
            { 1,              0,          1,         10,        1000,         1   },
            { 1,              10,         1,         10,        1000,         1   },
            { 1,              0,          1,         1000,      1000,         10  },
            { 1,              0,          10,        1000,      1000,         50  },
            { 1,              0,          1000,      10 ,       1000,         50  },
            { 1,              100,        1,         1000 * 4,  1000,         50  },
            { 100,            0,          1000,      1,         2,            1   },
            { 3,              0,          4,         5006,      1000,         50  },
            { 20,             0,          100,       4888,      1000,         100 },
            { 16,             10,         1000,      50 ,       1000,         100 },
            { 16,             10,         1000,      50 ,       657,          657 },
        };
    }

    @Test(dataProvider = "config")
    public void test(int iterations,
                     int delayMillis,
                     int numBuffers,
                     int bufferSize,
                     int maxBufferSize,
                     int minbufferSize) {
        for (long perRequestAmount : new long[] { 1L, MAX_VALUE })
            test(iterations,
                 delayMillis,
                 numBuffers,
                 bufferSize,
                 maxBufferSize,
                 minbufferSize,
                 perRequestAmount);
    }

    volatile boolean onNextThrew;

    BiConsumer<Flow.Subscriber<?>, ? super Throwable> onNextThrowHandler =
            (sub, ex) -> {
                onNextThrew = true;
                System.out.println("onNext threw " + ex);
                ex.printStackTrace();
    };

    public void test(int iterations,
                     int delayMillis,
                     int numBuffers,
                     int bufferSize,
                     int maxBufferSize,
                     int minBufferSize,
                     long requestAmount) {
        ExecutorService executor = Executors.newFixedThreadPool(1);
        try {
            out.printf("Iterations %d\n", iterations);
            for (int i=0; i<iterations; i++ ) {
                printStamp(START, "Iteration %d", i);
                try {
                    SubmissionPublisher<List<ByteBuffer>> publisher =
                            new SubmissionPublisher<>(executor,
                                                      1, // lock-step with the publisher, for now
                                                      onNextThrowHandler);
                    CompletableFuture<?> cf = sink(publisher,
                            delayMillis,
                            numBuffers * bufferSize,
                            requestAmount,
                            maxBufferSize,
                            minBufferSize);
                    source(publisher, numBuffers, bufferSize);
                    publisher.close();
                    cf.join();
                } finally {
                    printStamp(END, "Iteration %d\n", i);
                }
            }

            assertFalse(onNextThrew, "Unexpected onNextThrew, check output");

            out.println("OK");
        } finally {
            executor.shutdown();
        }
    }

    static long accumulatedDataSize(List<ByteBuffer> bufs) {
        return bufs.stream().mapToLong(ByteBuffer::remaining).sum();
    }

    /** Returns a new BB with its contents set to monotonically increasing
     * values, staring at the given start index and wrapping every 100. */
    static ByteBuffer allocateBuffer(int size, int startIdx) {
        ByteBuffer b = ByteBuffer.allocate(size);
        for (int i=0; i<size; i++)
            b.put((byte)((startIdx + i) % 100));
        b.position(0);
        return b;
    }

    static class TestSubscriber implements BodySubscriber<Integer> {
        final int delayMillis;
        final int bufferSize;
        final int expectedTotalSize;
        final long requestAmount;
        final CompletableFuture<Integer> completion;
        final Executor delayedExecutor;
        volatile Flow.Subscription subscription;

        TestSubscriber(int bufferSize,
                       int delayMillis,
                       int expectedTotalSize,
                       long requestAmount) {
            this.bufferSize = bufferSize;
            this.completion = new CompletableFuture<>();
            this.delayMillis = delayMillis;
            this.delayedExecutor = delayedExecutor(delayMillis, MILLISECONDS);
            this.expectedTotalSize = expectedTotalSize;
            this.requestAmount = requestAmount;
        }

        /**
         * Example of a factory method which would decorate a buffering
         * subscriber to create a new subscriber dependent on buffering capability.
         * <p>
         * The integer type parameter simulates the body just by counting the
         * number of bytes in the body.
         */
        static BodySubscriber<Integer> createSubscriber(int bufferSize,
                                                        int delay,
                                                        int expectedTotalSize,
                                                        long requestAmount) {
            TestSubscriber s = new TestSubscriber(bufferSize,
                                                delay,
                                                expectedTotalSize,
                                                requestAmount);
            return BodySubscribers.buffering(s, bufferSize);
        }

        private void requestMore() {
            subscription.request(requestAmount);
        }

        @Override
        public void onSubscribe(Subscription subscription) {
            assertNull(this.subscription);
            this.subscription = subscription;
            if (delayMillis > 0)
                delayedExecutor.execute(this::requestMore);
            else
                requestMore();
        }

        volatile int wrongSizes;
        volatile int totalBytesReceived;
        volatile int onNextInvocations;
        volatile long lastSeenSize = -1;
        volatile boolean noMoreOnNext; // false
        volatile int index; // 0
        volatile long count;

        @Override
        public void onNext(List<ByteBuffer> items) {
            try {
                long sz = accumulatedDataSize(items);
                boolean printStamp = delayMillis > 0
                        && requestAmount < Long.MAX_VALUE
                        && count % 20 == 0;
                if (printStamp) {
                    printStamp("stamp", "count=%d sz=%d accumulated=%d",
                            count, sz, (totalBytesReceived + sz));
                }
                count++;
                onNextInvocations++;
                assertNotEquals(sz, 0L, "Unexpected empty buffers");
                items.stream().forEach(b -> assertEquals(b.position(), 0));
                assertFalse(noMoreOnNext);

                if (sz != bufferSize) {
                    String msg = sz + ", should be less than bufferSize, " + bufferSize;
                    assertTrue(sz < bufferSize, msg);
                    assertTrue(lastSeenSize == -1 || lastSeenSize == bufferSize);
                    noMoreOnNext = true;
                    wrongSizes++;
                    printStamp("onNext",
                            "Possibly received last buffer: sz=%d, accumulated=%d, total=%d",
                            sz, totalBytesReceived, totalBytesReceived + sz);
                } else {
                    assertEquals(sz, bufferSize, "Expected to receive exactly bufferSize");
                }
                lastSeenSize = sz;

                // Ensure expected contents
                for (ByteBuffer b : items) {
                    while (b.hasRemaining()) {
                        assertEquals(b.get(), (byte) (index % 100));
                        index++;
                    }
                }

                totalBytesReceived += sz;
                assertEquals(totalBytesReceived, index);
                if (delayMillis > 0 && ((expectedTotalSize - totalBytesReceived) > bufferSize))
                    delayedExecutor.execute(this::requestMore);
                else
                    requestMore();
            } catch (Throwable t) {
                completion.completeExceptionally(t);
            }
        }

        @Override
        public void onError(Throwable throwable) {
            completion.completeExceptionally(throwable);
        }

        @Override
        public void onComplete() {
            if (wrongSizes > 1) { // allow just the final item to be smaller
                String msg = "Wrong sizes. Expected no more than 1. [" + this + "]";
                completion.completeExceptionally(new Throwable(msg));
            }
            if (totalBytesReceived != expectedTotalSize) {
                String msg = "Wrong number of bytes. [" + this + "]";
                completion.completeExceptionally(new Throwable(msg));
            } else {
                completion.complete(totalBytesReceived);
            }
        }

        @Override
        public CompletionStage<Integer> getBody() {
            return completion;
        }

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder();
            sb.append(super.toString());
            sb.append(", bufferSize=").append(bufferSize);
            sb.append(", onNextInvocations=").append(onNextInvocations);
            sb.append(", totalBytesReceived=").append(totalBytesReceived);
            sb.append(", expectedTotalSize=").append(expectedTotalSize);
            sb.append(", requestAmount=").append(requestAmount);
            sb.append(", lastSeenSize=").append(lastSeenSize);
            sb.append(", wrongSizes=").append(wrongSizes);
            sb.append(", index=").append(index);
            return sb.toString();
        }
    }

    /**
     * Publishes data, through the given publisher, using the main thread.
     *
     * Note: The executor supplied when creating the SubmissionPublisher provides
     * the threads for executing the Subscribers.
     *
     * @param publisher the publisher
     * @param numBuffers the number of buffers to send ( before splitting in two )
     * @param bufferSize the total size of the data to send ( before splitting in two )
     */
    static void source(SubmissionPublisher<List<ByteBuffer>> publisher,
                       int numBuffers,
                       int bufferSize) {
        printStamp("source","Publishing %d buffers of size %d each", numBuffers, bufferSize);
        int index = 0;
        for (int i=0; i<numBuffers; i++) {
            int chunkSize = random.nextInt(bufferSize);
            ByteBuffer buf1 = allocateBuffer(chunkSize, index);
            index += chunkSize;
            ByteBuffer buf2 = allocateBuffer(bufferSize - chunkSize, index);
            index += bufferSize - chunkSize;
            publisher.submit(List.of(buf1, buf2));
        }
        printStamp("source", "complete");
    }

    /**
     * Creates and subscribes Subscribers that receive data from the given
     * publisher.
     *
     * @param publisher the publisher
     * @param delayMillis time, in milliseconds, to delay the Subscription
     *                    requesting more bytes ( for simulating slow consumption )
     * @param expectedTotalSize the total number of bytes expected to be received
     *                          by the subscribers
     * @return a CompletableFuture which completes when the subscription is complete
     */
    static CompletableFuture<?> sink(SubmissionPublisher<List<ByteBuffer>> publisher,
                                     int delayMillis,
                                     int expectedTotalSize,
                                     long requestAmount,
                                     int maxBufferSize,
                                     int minBufferSize) {
        int bufferSize = chooseBufferSize(maxBufferSize,
                                          minBufferSize,
                                          delayMillis,
                                          expectedTotalSize,
                                          requestAmount);
        assert bufferSize > 0;
        assert bufferSize >= minBufferSize;
        assert bufferSize <= maxBufferSize;
        BodySubscriber<Integer> sub = TestSubscriber.createSubscriber(bufferSize,
                                                                    delayMillis,
                                                                    expectedTotalSize,
                                                                    requestAmount);
        publisher.subscribe(sub);
        printStamp("sink","Subscriber reads data with buffer size: %d", bufferSize);
        out.printf("Subscription delay is %d msec\n", delayMillis);
        long delay = (((long)delayMillis * expectedTotalSize) / bufferSize) / requestAmount;
        out.printf("Minimum total delay is %d sec %d ms\n", delay / 1000, delay % 1000);
        out.printf("Request amount is %d items\n", requestAmount);
        return sub.getBody().toCompletableFuture();
    }

    static int chooseBufferSize(int maxBufferSize,
                                int minBufferSize,
                                int delaysMillis,
                                int expectedTotalSize,
                                long requestAmount) {
        assert minBufferSize > 0 && maxBufferSize > 0 && requestAmount > 0;
        int bufferSize = maxBufferSize == minBufferSize ? maxBufferSize :
                (random.nextInt(maxBufferSize - minBufferSize)
                    + minBufferSize);
        if (requestAmount == Long.MAX_VALUE) return bufferSize;
        long minDelay = (((long)delaysMillis * expectedTotalSize) / maxBufferSize)
                / requestAmount;
        long maxDelay = (((long)delaysMillis * expectedTotalSize) / minBufferSize)
                / requestAmount;
        // if the maximum delay is < 10s just take a random number between min and max.
        if (maxDelay <= LOWER_THRESHOLD) {
            return bufferSize;
        }
        // if minimum delay is greater than 20s then print a warning and use max buffer.
        if (minDelay >= UPPER_THRESHOLD) {
            System.out.println("Warning: minimum delay is "
                    + minDelay/1000 + "sec " + minDelay%1000 + "ms");
            System.err.println("Warning: minimum delay is "
                    + minDelay/1000 + "sec " + minDelay%1000 + "ms");
            return maxBufferSize;
        }
        // maxDelay could be anything, but minDelay is below the UPPER_THRESHOLD
        // try to pick up a buffer size that keeps the delay below the
        // UPPER_THRESHOLD
        while (minBufferSize < maxBufferSize) {
            bufferSize = random.nextInt(maxBufferSize - minBufferSize)
                    + minBufferSize;
            long delay = (((long)delaysMillis * expectedTotalSize) / bufferSize)
                    / requestAmount;
            if (delay < UPPER_THRESHOLD) return bufferSize;
            minBufferSize++;
        }
        return minBufferSize;
    }

    // ---

    /* Main entry point for standalone testing of the main functional test. */
    public static void main(String... args) {
        BufferingSubscriberTest t = new BufferingSubscriberTest();
        for (Object[] objs : t.config())
            t.test((int)objs[0], (int)objs[1], (int)objs[2], (int)objs[3], (int)objs[4], (int)objs[5]);
    }
}