/*

 * Copyright (c) 2016, 2019, 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.  Oracle designates this

 * particular file as subject to the "Classpath" exception as provided

 * by Oracle in the LICENSE file that accompanied this code.

 *

 * 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.

 */

package jdk.internal.net.http;

import java.io.BufferedReader;

import java.io.FilePermission;

import java.io.IOException;

import java.io.InputStream;

import java.io.InputStreamReader;

import java.nio.ByteBuffer;

import java.nio.channels.FileChannel;

import java.nio.charset.Charset;

import java.nio.file.OpenOption;

import java.nio.file.Path;

import java.security.AccessController;

import java.security.PrivilegedActionException;

import java.security.PrivilegedExceptionAction;

import java.util.ArrayList;

import java.util.Iterator;

import java.util.List;

import java.util.Objects;

import java.util.Optional;

import java.util.concurrent.ArrayBlockingQueue;

import java.util.concurrent.BlockingQueue;

import java.util.concurrent.CompletableFuture;

import java.util.concurrent.CompletionStage;

import java.util.concurrent.Executor;

import java.util.concurrent.Flow;

import java.util.concurrent.Flow.Subscriber;

import java.util.concurrent.Flow.Subscription;

import java.util.concurrent.atomic.AtomicBoolean;

import java.util.concurrent.atomic.AtomicReference;

import java.util.function.Consumer;

import java.util.function.Function;

import java.util.stream.Stream;

import java.net.http.HttpResponse.BodySubscriber;

import jdk.internal.net.http.common.Log;

import jdk.internal.net.http.common.Logger;

import jdk.internal.net.http.common.MinimalFuture;

import jdk.internal.net.http.common.Utils;

import static java.nio.charset.StandardCharsets.UTF_8;

public class ResponseSubscribers {

    /**

     * This interface is used by our BodySubscriber implementations to

     * declare whether calling getBody() inline is safe, or whether

     * it needs to be called asynchronously in an executor thread.

     * Calling getBody() inline is usually safe except when it

     * might block - which can be the case if the BodySubscriber

     * is provided by custom code, or if it uses a finisher that

     * might be called and might block before the last bit is

     * received (for instance, if a mapping subscriber is used with

     * a mapper function that maps an InputStream to a GZIPInputStream,

     * as the the constructor of GZIPInputStream calls read()).

     * @param <T> The response type.

     */

    public interface TrustedSubscriber<T> extends BodySubscriber<T> {

        /**

         * Returns true if getBody() should be called asynchronously.

         * @implSpec The default implementation of this method returns

         *           false.

         * @return true if getBody() should be called asynchronously.

         */

        default boolean needsExecutor() { return false;}

        /**

         * Returns true if calling {@code bs::getBody} might block

         * and requires an executor.

         *

         * @implNote

         * In particular this method returns

         * true if {@code bs} is not a {@code TrustedSubscriber}.

         * If it is a {@code TrustedSubscriber}, it returns

         * {@code ((TrustedSubscriber) bs).needsExecutor()}.

         *

         * @param bs A BodySubscriber.

         * @return true if calling {@code bs::getBody} requires using

         *         an executor.

         */

        static boolean needsExecutor(BodySubscriber<?> bs) {

            if (bs instanceof TrustedSubscriber) {

                return ((TrustedSubscriber) bs).needsExecutor();

            } else return true;

        }

    }

    public static class ConsumerSubscriber implements TrustedSubscriber<Void> {

        private final Consumer<Optional<byte[]>> consumer;

        private Flow.Subscription subscription;

        private final CompletableFuture<Void> result = new MinimalFuture<>();

        private final AtomicBoolean subscribed = new AtomicBoolean();

        public ConsumerSubscriber(Consumer<Optional<byte[]>> consumer) {

            this.consumer = Objects.requireNonNull(consumer);

        }

        @Override

        public CompletionStage<Void> getBody() {

            return result;

        }

        @Override

        public void onSubscribe(Flow.Subscription subscription) {

            Objects.requireNonNull(subscription);

            if (!subscribed.compareAndSet(false, true)) {

                subscription.cancel();

            } else {

                this.subscription = subscription;

                subscription.request(1);

            }

        }

        @Override

        public void onNext(List<ByteBuffer> items) {

            Objects.requireNonNull(items);

            for (ByteBuffer item : items) {

                byte[] buf = new byte[item.remaining()];

                item.get(buf);

                consumer.accept(Optional.of(buf));

            }

            subscription.request(1);

        }

        @Override

        public void onError(Throwable throwable) {

            Objects.requireNonNull(throwable);

            result.completeExceptionally(throwable);

        }

        @Override

        public void onComplete() {

            consumer.accept(Optional.empty());

            result.complete(null);

        }

    }

    /**

     * A Subscriber that writes the flow of data to a given file.

     *

     * Privileged actions are performed within a limited doPrivileged that only

     * asserts the specific, write, file permissions that were checked during

     * the construction of this PathSubscriber.

     */

    public static class PathSubscriber implements TrustedSubscriber<Path> {

        private static final FilePermission[] EMPTY_FILE_PERMISSIONS = new FilePermission[0];

        private final Path file;

        private final OpenOption[] options;

        private final FilePermission[] filePermissions;

        private final CompletableFuture<Path> result = new MinimalFuture<>();

        private final AtomicBoolean subscribed = new AtomicBoolean();

        private volatile Flow.Subscription subscription;

        private volatile FileChannel out;

        private static final String pathForSecurityCheck(Path path) {

            return path.toFile().getPath();

        }

        /**

         * Factory for creating PathSubscriber.

         *

         * Permission checks are performed here before construction of the

         * PathSubscriber. Permission checking and construction are deliberately

         * and tightly co-located.

         */

        public static PathSubscriber create(Path file,

                                            List<OpenOption> options) {

            FilePermission filePermission = null;

            SecurityManager sm = System.getSecurityManager();

            if (sm != null) {

                String fn = pathForSecurityCheck(file);

                FilePermission writePermission = new FilePermission(fn, "write");

                sm.checkPermission(writePermission);

                filePermission = writePermission;

            }

            return new PathSubscriber(file, options, filePermission);

        }

        // pp so handler implementations in the same package can construct

        /*package-private*/ PathSubscriber(Path file,

                                           List<OpenOption> options,

                                           FilePermission... filePermissions) {

            this.file = file;

            this.options = options.stream().toArray(OpenOption[]::new);

            this.filePermissions =

                    filePermissions == null ? EMPTY_FILE_PERMISSIONS : filePermissions;

        }

        @Override

        public void onSubscribe(Flow.Subscription subscription) {

            Objects.requireNonNull(subscription);

            if (!subscribed.compareAndSet(false, true)) {

                subscription.cancel();

                return;

            }

            this.subscription = subscription;

            if (System.getSecurityManager() == null) {

                try {

                    out = FileChannel.open(file, options);

                } catch (IOException ioe) {

                    result.completeExceptionally(ioe);

                    return;

                }

            } else {

                try {

                    PrivilegedExceptionAction<FileChannel> pa =

                            () -> FileChannel.open(file, options);

                    out = AccessController.doPrivileged(pa, null, filePermissions);

                } catch (PrivilegedActionException pae) {

                    Throwable t = pae.getCause() != null ? pae.getCause() : pae;

                    result.completeExceptionally(t);

                    subscription.cancel();

                    return;

                }

            }

            subscription.request(1);

        }

        @Override

        public void onNext(List<ByteBuffer> items) {

            try {

                out.write(items.toArray(Utils.EMPTY_BB_ARRAY));

            } catch (IOException ex) {

                Utils.close(out);

                subscription.cancel();

                result.completeExceptionally(ex);

            }

            subscription.request(1);

        }

        @Override

        public void onError(Throwable e) {

            result.completeExceptionally(e);

            Utils.close(out);

        }

        @Override

        public void onComplete() {

            Utils.close(out);

            result.complete(file);

        }

        @Override

        public CompletionStage<Path> getBody() {

            return result;

        }

    }

    public static class ByteArraySubscriber<T> implements TrustedSubscriber<T> {

        private final Function<byte[], T> finisher;

        private final CompletableFuture<T> result = new MinimalFuture<>();

        private final List<ByteBuffer> received = new ArrayList<>();

        private volatile Flow.Subscription subscription;

        public ByteArraySubscriber(Function<byte[],T> finisher) {

            this.finisher = finisher;

        }

        @Override

        public void onSubscribe(Flow.Subscription subscription) {

            if (this.subscription != null) {

                subscription.cancel();

                return;

            }

            this.subscription = subscription;

            // We can handle whatever you've got

            subscription.request(Long.MAX_VALUE);

        }

        @Override

        public void onNext(List<ByteBuffer> items) {

            // incoming buffers are allocated by http client internally,

            // and won't be used anywhere except this place.

            // So it's free simply to store them for further processing.

            assert Utils.hasRemaining(items);

            received.addAll(items);

        }

        @Override

        public void onError(Throwable throwable) {

            received.clear();

            result.completeExceptionally(throwable);

        }

        static private byte[] join(List<ByteBuffer> bytes) {

            int size = Utils.remaining(bytes, Integer.MAX_VALUE);

            byte[] res = new byte[size];

            int from = 0;

            for (ByteBuffer b : bytes) {

                int l = b.remaining();

                b.get(res, from, l);

                from += l;

            }

            return res;

        }

        @Override

        public void onComplete() {

            try {

                result.complete(finisher.apply(join(received)));

                received.clear();

            } catch (IllegalArgumentException e) {

                result.completeExceptionally(e);

            }

        }

        @Override

        public CompletionStage<T> getBody() {

            return result;

        }

    }

    /**

     * An InputStream built on top of the Flow API.

     */

    public static class HttpResponseInputStream extends InputStream

        implements TrustedSubscriber<InputStream>

    {

        final static int MAX_BUFFERS_IN_QUEUE = 1;  // lock-step with the producer

        // An immutable ByteBuffer sentinel to mark that the last byte was received.

        private static final ByteBuffer LAST_BUFFER = ByteBuffer.wrap(new byte[0]);

        private static final List<ByteBuffer> LAST_LIST = List.of(LAST_BUFFER);

        private static final Logger debug =

                Utils.getDebugLogger("HttpResponseInputStream"::toString, Utils.DEBUG);

        // A queue of yet unprocessed ByteBuffers received from the flow API.

        private final BlockingQueue<List<ByteBuffer>> buffers;

        private volatile Flow.Subscription subscription;

        private volatile boolean closed;

        private volatile Throwable failed;

        private volatile Iterator<ByteBuffer> currentListItr;

        private volatile ByteBuffer currentBuffer;

        private final AtomicBoolean subscribed = new AtomicBoolean();

        public HttpResponseInputStream() {

            this(MAX_BUFFERS_IN_QUEUE);

        }

        HttpResponseInputStream(int maxBuffers) {

            int capacity = (maxBuffers <= 0 ? MAX_BUFFERS_IN_QUEUE : maxBuffers);

            // 1 additional slot needed for LAST_LIST added by onComplete

            this.buffers = new ArrayBlockingQueue<>(capacity + 1);

        }

        @Override

        public CompletionStage<InputStream> getBody() {

            // Returns the stream immediately, before the

            // response body is received.

            // This makes it possible for sendAsync().get().body()

            // to complete before the response body is received.

            return CompletableFuture.completedStage(this);

        }

        // Returns the current byte buffer to read from.

        // If the current buffer has no remaining data, this method will take the

        // next buffer from the buffers queue, possibly blocking until

        // a new buffer is made available through the Flow API, or the

        // end of the flow has been reached.

        private ByteBuffer current() throws IOException {

            while (currentBuffer == null || !currentBuffer.hasRemaining()) {

                // Check whether the stream is closed or exhausted

                if (closed || failed != null) {

                    throw new IOException("closed", failed);

                }

                if (currentBuffer == LAST_BUFFER) break;

                try {

                    if (currentListItr == null || !currentListItr.hasNext()) {

                        // Take a new list of buffers from the queue, blocking

                        // if none is available yet...

                        if (debug.on()) debug.log("Taking list of Buffers");

                        List<ByteBuffer> lb = buffers.take();

                        currentListItr = lb.iterator();

                        if (debug.on()) debug.log("List of Buffers Taken");

                        // Check whether an exception was encountered upstream

                        if (closed || failed != null)

                            throw new IOException("closed", failed);

                        // Check whether we're done.

                        if (lb == LAST_LIST) {

                            currentListItr = null;

                            currentBuffer = LAST_BUFFER;

                            break;

                        }

                        // Request another upstream item ( list of buffers )

                        Flow.Subscription s = subscription;

                        if (s != null) {

                            if (debug.on()) debug.log("Increased demand by 1");

                            s.request(1);

                        }

                        assert currentListItr != null;

                        if (lb.isEmpty()) continue;

                    }

                    assert currentListItr != null;

                    assert currentListItr.hasNext();

                    if (debug.on()) debug.log("Next Buffer");

                    currentBuffer = currentListItr.next();

                } catch (InterruptedException ex) {

                    // continue

                }

            }

            assert currentBuffer == LAST_BUFFER || currentBuffer.hasRemaining();

            return currentBuffer;

        }

        @Override

        public int read(byte[] bytes, int off, int len) throws IOException {

            Objects.checkFromIndexSize(off, len, bytes.length);

            if (len == 0) {

                return 0;

            }

            // get the buffer to read from, possibly blocking if

            // none is available

            ByteBuffer buffer;

            if ((buffer = current()) == LAST_BUFFER) return -1;

            // don't attempt to read more than what is available

            // in the current buffer.

            int read = Math.min(buffer.remaining(), len);

            assert read > 0 && read <= buffer.remaining();

            // buffer.get() will do the boundary check for us.

            buffer.get(bytes, off, read);

            return read;

        }

        @Override

        public int read() throws IOException {

            ByteBuffer buffer;

            if ((buffer = current()) == LAST_BUFFER) return -1;

            return buffer.get() & 0xFF;

        }

        @Override

        public int available() throws IOException {

            // best effort: returns the number of remaining bytes in

            // the current buffer if any, or 1 if the current buffer

            // is null or empty but the queue or current buffer list

            // are not empty. Returns 0 otherwise.

            if (closed) return 0;

            int available = 0;

            ByteBuffer current = currentBuffer;

            if (current == LAST_BUFFER) return 0;

            if (current != null) available = current.remaining();

            if (available != 0) return available;

            Iterator<?> iterator = currentListItr;

            if (iterator != null && iterator.hasNext()) return 1;

            if (buffers.isEmpty()) return 0;

            return 1;

        }

        @Override

        public void onSubscribe(Flow.Subscription s) {

            Objects.requireNonNull(s);

            try {

                if (!subscribed.compareAndSet(false, true)) {

                    s.cancel();

                } else {

                    // check whether the stream is already closed.

                    // if so, we should cancel the subscription

                    // immediately.

                    boolean closed;

                    synchronized (this) {

                        closed = this.closed;

                        if (!closed) {

                            this.subscription = s;

                        }

                    }

                    if (closed) {

                        s.cancel();

                        return;

                    }

                    assert buffers.remainingCapacity() > 1; // should contain at least 2

                    if (debug.on())

                        debug.log("onSubscribe: requesting "

                                  + Math.max(1, buffers.remainingCapacity() - 1));

                    s.request(Math.max(1, buffers.remainingCapacity() - 1));

                }

            } catch (Throwable t) {