/*

 * Copyright (c) 2015, 2017, 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.incubator.http.internal.websocket;

import java.io.IOException;

import java.net.ProtocolException;

import java.net.URI;

import java.nio.ByteBuffer;

import java.util.Queue;

import java.util.concurrent.CompletableFuture;

import java.util.concurrent.CompletionStage;

import java.util.concurrent.ConcurrentLinkedQueue;

import java.util.concurrent.atomic.AtomicBoolean;

import java.util.function.Consumer;

import java.util.function.Function;

import jdk.incubator.http.WebSocket;

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

import jdk.incubator.http.internal.common.Pair;

import jdk.incubator.http.internal.common.SequentialScheduler;

import jdk.incubator.http.internal.common.SequentialScheduler.DeferredCompleter;

import jdk.incubator.http.internal.websocket.OpeningHandshake.Result;

import jdk.incubator.http.internal.websocket.OutgoingMessage.Binary;

import jdk.incubator.http.internal.websocket.OutgoingMessage.Close;

import jdk.incubator.http.internal.websocket.OutgoingMessage.Context;

import jdk.incubator.http.internal.websocket.OutgoingMessage.Ping;

import jdk.incubator.http.internal.websocket.OutgoingMessage.Pong;

import jdk.incubator.http.internal.websocket.OutgoingMessage.Text;

import static java.util.Objects.requireNonNull;

import static java.util.concurrent.CompletableFuture.failedFuture;

import static jdk.incubator.http.internal.common.Pair.pair;

import static jdk.incubator.http.internal.websocket.StatusCodes.CLOSED_ABNORMALLY;

import static jdk.incubator.http.internal.websocket.StatusCodes.NO_STATUS_CODE;

import static jdk.incubator.http.internal.websocket.StatusCodes.isLegalToSendFromClient;

/*

 * A WebSocket client.

 */

final class WebSocketImpl implements WebSocket {

    private final URI uri;

    private final String subprotocol;

    private final RawChannel channel;

    private final Listener listener;

    /*

     * Whether or not Listener.onClose or Listener.onError has been already

     * invoked. We keep track of this since only one of these methods is invoked

     * and it is invoked at most once.

     */

    private boolean lastMethodInvoked;

    private final AtomicBoolean outstandingSend = new AtomicBoolean();

    private final SequentialScheduler sendScheduler;

    private final Queue<Pair<OutgoingMessage, CompletableFuture<WebSocket>>>

            queue = new ConcurrentLinkedQueue<>();

    private final Context context = new OutgoingMessage.Context();

    private final Transmitter transmitter;

    private final Receiver receiver;

    /*

     * Whether or not the WebSocket has been closed. When a WebSocket has been

     * closed it means that no further messages can be sent or received.

     * A closure can be triggered by:

     *

     *   1. abort()

     *   2. "Failing the WebSocket Connection" (i.e. a fatal error)

     *   3. Completion of the Closing handshake

     */

    private final AtomicBoolean closed = new AtomicBoolean();

    /*

     * This lock is enforcing sequential ordering of invocations to listener's

     * methods. It is supposed to be uncontended. The only contention that can

     * happen is when onOpen, an asynchronous onError (not related to reading

     * from the channel, e.g. an error from automatic Pong reply) or onClose

     * (related to abort) happens. Since all of the above are one-shot actions,

     * the said contention is insignificant.

     */

    private final Object lock = new Object();

    private final CompletableFuture<?> closeReceived = new CompletableFuture<>();

    private final CompletableFuture<?> closeSent = new CompletableFuture<>();

    static CompletableFuture<WebSocket> newInstanceAsync(BuilderImpl b) {

        Function<Result, WebSocket> newWebSocket = r -> {

            WebSocketImpl ws = new WebSocketImpl(b.getUri(),

                                                 r.subprotocol,

                                                 r.channel,

                                                 b.getListener());

            // The order of calls might cause a subtle effects, like CF will be

            // returned from the buildAsync _after_ onOpen has been signalled.

            // This means if onOpen is lengthy, it might cause some problems.

            ws.signalOpen();

            return ws;

        };

        OpeningHandshake h;

        try {

        } catch (IllegalArgumentException e) {

            return failedFuture(e);

        }

        return h.send().thenApply(newWebSocket);

    }

    WebSocketImpl(URI uri,

                  String subprotocol,

                  RawChannel channel,

                  Listener listener)

    {

        this.uri = requireNonNull(uri);

        this.subprotocol = requireNonNull(subprotocol);

        this.channel = requireNonNull(channel);

        this.listener = requireNonNull(listener);

        this.transmitter = new Transmitter(channel);

        this.receiver = new Receiver(messageConsumerOf(listener), channel);

        this.sendScheduler = new SequentialScheduler(new SendFirstTask());

        // Set up the Closing Handshake action

        CompletableFuture.allOf(closeReceived, closeSent)

                .whenComplete((result, error) -> {

                    try {

                        channel.close();

                    } catch (IOException e) {

                        Log.logError(e);

                    } finally {

                        closed.set(true);

                    }

                });

    }

    /*

     * This initialisation is outside of the constructor for the sake of

     * safe publication.

     */

    private void signalOpen() {

        synchronized (lock) {

            // TODO: might hold lock longer than needed causing prolonged

            // contention? substitute lock for ConcurrentLinkedQueue<Runnable>?

            try {

                listener.onOpen(this);

            } catch (Exception e) {

                signalError(e);

            }

        }

    }

    private void signalError(Throwable error) {

        synchronized (lock) {

            if (lastMethodInvoked) {

                Log.logError(error);

            } else {

                lastMethodInvoked = true;

                receiver.close();

                try {

                    listener.onError(this, error);

                } catch (Exception e) {

                    Log.logError(e);

                }

            }

        }

    }

    /*

     * Processes a Close event that came from the channel. Invoked at most once.

     */

    private void processClose(int statusCode, String reason) {

        receiver.close();

        try {

            channel.shutdownInput();

        } catch (IOException e) {

            Log.logError(e);

        }

        boolean alreadyCompleted = !closeReceived.complete(null);

        if (alreadyCompleted) {

            // This CF is supposed to be completed only once, the first time a

            // Close message is received. No further messages are pulled from

            // the socket.

            throw new InternalError();

        }

        int code;

        if (statusCode == NO_STATUS_CODE || statusCode == CLOSED_ABNORMALLY) {

            code = NORMAL_CLOSURE;

        } else {

            code = statusCode;

        }

        CompletionStage<?> readyToClose = signalClose(statusCode, reason);

        if (readyToClose == null) {

            readyToClose = CompletableFuture.completedFuture(null);

        }

        readyToClose.whenComplete((r, error) -> {

            enqueueClose(new Close(code, ""))

                    .whenComplete((r1, error1) -> {

                        if (error1 != null) {

                            Log.logError(error1);

                        }

                    });

        });

    }

    /*

     * Signals a Close event (might not correspond to anything happened on the

     * channel, e.g. `abort()`).

     */

    private CompletionStage<?> signalClose(int statusCode, String reason) {

        synchronized (lock) {

            if (lastMethodInvoked) {

                Log.logTrace("Close: {0}, ''{1}''", statusCode, reason);

            } else {

                lastMethodInvoked = true;

                receiver.close();

                try {

                    return listener.onClose(this, statusCode, reason);

                } catch (Exception e) {

                    Log.logError(e);

                }

            }

        }

        return null;

    }

    @Override

    public CompletableFuture<WebSocket> sendText(CharSequence message,

                                                 boolean isLast)

    {

        return enqueueExclusively(new Text(message, isLast));

    }

    @Override

    public CompletableFuture<WebSocket> sendBinary(ByteBuffer message,

                                                   boolean isLast)

    {

        return enqueueExclusively(new Binary(message, isLast));

    }

    @Override

    public CompletableFuture<WebSocket> sendPing(ByteBuffer message) {

        return enqueueExclusively(new Ping(message));

    }

    @Override

    public CompletableFuture<WebSocket> sendPong(ByteBuffer message) {

        return enqueueExclusively(new Pong(message));

    }

    @Override

    public CompletableFuture<WebSocket> sendClose(int statusCode,

                                                  String reason) {

        if (!isLegalToSendFromClient(statusCode)) {

            return failedFuture(

                    new IllegalArgumentException("statusCode: " + statusCode));

        }

        Close msg;

        try {

            msg = new Close(statusCode, reason);

        } catch (IllegalArgumentException e) {

            return failedFuture(e);

        }

        return enqueueClose(msg);

    }

    /*

     * Sends a Close message with the given contents and then shuts down the

     * channel for writing since no more messages are expected to be sent after

     * this. Invoked at most once.

     */

    private CompletableFuture<WebSocket> enqueueClose(Close m) {

        return enqueue(m).whenComplete((r, error) -> {

            try {

                channel.shutdownOutput();

            } catch (IOException e) {

                Log.logError(e);

            }

            boolean alreadyCompleted = !closeSent.complete(null);

            if (alreadyCompleted) {

                // Shouldn't happen as this callback must run at most once

                throw new InternalError();

            }

        });

    }

    /*

     * Accepts the given message into the outgoing queue in a mutually-exclusive

     * fashion in respect to other messages accepted through this method. No

     * further messages will be accepted until the returned CompletableFuture

     * completes. This method is used to enforce "one outstanding send

     * operation" policy.

     */

    private CompletableFuture<WebSocket> enqueueExclusively(OutgoingMessage m)

    {

        if (closed.get()) {

            return failedFuture(new IllegalStateException("Closed"));

        }

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

            return failedFuture(new IllegalStateException("Outstanding send"));

        }

        return enqueue(m).whenComplete((r, e) -> outstandingSend.set(false));

    }

    private CompletableFuture<WebSocket> enqueue(OutgoingMessage m) {

        CompletableFuture<WebSocket> cf = new CompletableFuture<>();

        boolean added = queue.add(pair(m, cf));

        if (!added) {

            // The queue is supposed to be unbounded

            throw new InternalError();

        }

        sendScheduler.runOrSchedule();

        return cf;

    }

    /*

     * This is the main sending task. It may be run in different threads,

     * but never concurrently.

     */

    private class SendFirstTask implements SequentialScheduler.RestartableTask {

        @Override

        public void run (DeferredCompleter taskCompleter){

            Pair<OutgoingMessage, CompletableFuture<WebSocket>> p = queue.poll();

            if (p == null) {

                taskCompleter.complete();

                return;

            }

            OutgoingMessage message = p.first;

            CompletableFuture<WebSocket> cf = p.second;

            try {

                message.contextualize(context);

                Consumer<Exception> h = e -> {

                    if (e == null) {

                        cf.complete(WebSocketImpl.this);

                    } else {

                        cf.completeExceptionally(e);

                    }

                    sendScheduler.runOrSchedule();

                    taskCompleter.complete();

                };

                transmitter.send(message, h);

            } catch (Exception t) {

                cf.completeExceptionally(t);

            }

        }

    }

    @Override

    public void request(long n) {

        receiver.request(n);

    }

    @Override

    public String getSubprotocol() {

        return subprotocol;

    }

    @Override

    public boolean isClosed() {

        return closed.get();

    }

    @Override

    public void abort() throws IOException {

        try {

            channel.close();

        } finally {

            closed.set(true);

            signalClose(CLOSED_ABNORMALLY, "");

        }

    }

    @Override

    public String toString() {

        return super.toString()

                + "[" + (closed.get() ? "CLOSED" : "OPEN") + "]: " + uri

                + (!subprotocol.isEmpty() ? ", subprotocol=" + subprotocol : "");

    }

    private MessageStreamConsumer messageConsumerOf(Listener listener) {

        // Synchronization performed here in every method is not for the sake of

        // ordering invocations to this consumer, after all they are naturally

        // ordered in the channel. The reason is to avoid an interference with

        // any unrelated to the channel calls to onOpen, onClose and onError.

        return new MessageStreamConsumer() {

            @Override

            public void onText(MessagePart part, CharSequence data) {

                receiver.acknowledge();

                synchronized (WebSocketImpl.this.lock) {

                    try {

                        listener.onText(WebSocketImpl.this, data, part);

                    } catch (Exception e) {

                        signalError(e);

                    }

                }

            }

            @Override

            public void onBinary(MessagePart part, ByteBuffer data) {

                receiver.acknowledge();

                synchronized (WebSocketImpl.this.lock) {

                    try {