/*

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

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

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.SequentialScheduler;

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

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

import java.io.IOException;

import java.lang.ref.Reference;

import java.net.ProtocolException;

import java.net.URI;

import java.net.http.WebSocket;

import java.nio.ByteBuffer;

import java.nio.CharBuffer;

import java.nio.charset.CharacterCodingException;

import java.nio.charset.CharsetEncoder;

import java.nio.charset.CodingErrorAction;

import java.nio.charset.StandardCharsets;

import java.util.Objects;

import java.util.concurrent.CompletableFuture;

import java.util.concurrent.CompletionStage;

import java.util.concurrent.atomic.AtomicBoolean;

import java.util.concurrent.atomic.AtomicLong;

import java.util.concurrent.atomic.AtomicReference;

import java.util.function.BiConsumer;

import java.util.function.Function;

import static java.util.Objects.requireNonNull;

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

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

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

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

import static jdk.internal.net.http.websocket.WebSocketImpl.State.BINARY;

import static jdk.internal.net.http.websocket.WebSocketImpl.State.CLOSE;

import static jdk.internal.net.http.websocket.WebSocketImpl.State.ERROR;

import static jdk.internal.net.http.websocket.WebSocketImpl.State.IDLE;

import static jdk.internal.net.http.websocket.WebSocketImpl.State.OPEN;

import static jdk.internal.net.http.websocket.WebSocketImpl.State.PING;

import static jdk.internal.net.http.websocket.WebSocketImpl.State.PONG;

import static jdk.internal.net.http.websocket.WebSocketImpl.State.TEXT;

import static jdk.internal.net.http.websocket.WebSocketImpl.State.WAITING;

/*

 * A WebSocket client.

 */

public final class WebSocketImpl implements WebSocket {

    private static final Logger debug =

            Utils.getWebSocketLogger("[WebSocket]"::toString, Utils.DEBUG_WS);

    private final AtomicLong sendCounter = new AtomicLong();

    private final AtomicLong receiveCounter = new AtomicLong();

    enum State {

        OPEN,

        IDLE,

        WAITING,

        TEXT,

        BINARY,

        PING,

        PONG,

        CLOSE,

        ERROR

    }

    private final AtomicReference<ByteBuffer> lastAutomaticPong = new AtomicReference<>();

    private final MinimalFuture<WebSocket> DONE = MinimalFuture.completedFuture(this);

    private volatile boolean inputClosed;

    private final AtomicBoolean outputClosed = new AtomicBoolean();

    private final AtomicReference<State> state = new AtomicReference<>(OPEN);

    /* Components of calls to Listener's methods */

    private boolean last;

    private ByteBuffer binaryData;

    private CharSequence text;

    private int statusCode;

    private String reason;

    private final AtomicReference<Throwable> error = new AtomicReference<>();

    private final URI uri;

    private final String subprotocol;

    private final Listener listener;

    private final AtomicBoolean pendingTextOrBinary = new AtomicBoolean();

    private final AtomicBoolean pendingPingOrPong = new AtomicBoolean();

    private final Transport transport;

    private final SequentialScheduler receiveScheduler

            = new SequentialScheduler(new ReceiveTask());

    private final Demand demand = new Demand();

    public static CompletableFuture<WebSocket> newInstanceAsync(BuilderImpl b) {

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

            WebSocket ws = newInstance(b.getUri(),

                                       r.subprotocol,

                                       b.getListener(),

                                       r.transport);

            // Make sure we don't release the builder until this lambda

            // has been executed. The builder has a strong reference to

            // the HttpClientFacade, and we want to keep that live until

            // after the raw channel is created and passed to WebSocketImpl.

            Reference.reachabilityFence(b);

            return ws;

        };

        OpeningHandshake h;

        try {

            h = new OpeningHandshake(b);

        } catch (Throwable e) {

            return failedFuture(e);

        }

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

    }

    /* Exposed for testing purposes */

    static WebSocketImpl newInstance(URI uri,

                                     String subprotocol,

                                     Listener listener,

                                     TransportFactory transport) {

        WebSocketImpl ws = new WebSocketImpl(uri, subprotocol, listener, transport);

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

        // safe publication of WebSocketImpl.this

        ws.signalOpen();

        return ws;

    }

    private WebSocketImpl(URI uri,

                          String subprotocol,

                          Listener listener,

                          TransportFactory transportFactory) {

        this.uri = requireNonNull(uri);

        this.subprotocol = requireNonNull(subprotocol);

        this.listener = requireNonNull(listener);

        // Why 6? 1 sendPing/sendPong + 1 sendText/sendBinary + 1 Close +

        // 2 automatic Ping replies + 1 automatic Close = 6 messages

        // Why 2 automatic Pong replies? One is being sent, but the byte buffer

        // has been set to null, another just has been added.

        this.transport = transportFactory.createTransport(new MessageQueue(6),

                new SignallingMessageConsumer());

    }

    // FIXME: add to action handling of errors -> signalError()

    @Override

    public CompletableFuture<WebSocket> sendText(CharSequence message,

                                                 boolean last) {

        Objects.requireNonNull(message);

        long id = 0;

        if (debug.on()) {

            id = sendCounter.incrementAndGet();

            debug.log("enter send text %s payload length=%s last=%s",

                      id, message.length(), last);

        }

        CompletableFuture<WebSocket> result;

        if (!setPendingTextOrBinary()) {

            result = failedFuture(new IllegalStateException("Send pending"));

        } else {

            result = transport.sendText(message, last, this,

                                        (r, e) -> clearPendingTextOrBinary());

        }

        if (debug.on()) {

            debug.log("exit send text %s returned %s", id, result);

        }

        return replaceNull(result);

    }

    @Override

    public CompletableFuture<WebSocket> sendBinary(ByteBuffer message,

                                                   boolean last) {

        Objects.requireNonNull(message);

        long id = 0;

        if (debug.on()) {

            id = sendCounter.incrementAndGet();

            debug.log("enter send binary %s payload=%s last=%s",

                      id, message, last);

        }

        CompletableFuture<WebSocket> result;

        if (!setPendingTextOrBinary()) {

            result = failedFuture(new IllegalStateException("Send pending"));

        } else {

            result = transport.sendBinary(message, last, this,

                                          (r, e) -> clearPendingTextOrBinary());

        }

        if (debug.on()) {

            debug.log("exit send binary %s returned %s", id, result);

        }

        return replaceNull(result);

    }

    private void clearPendingTextOrBinary() {

        pendingTextOrBinary.set(false);

    }

    private boolean setPendingTextOrBinary() {

        return pendingTextOrBinary.compareAndSet(false, true);

    }

    private CompletableFuture<WebSocket> replaceNull(

            CompletableFuture<WebSocket> cf)

    {

        if (cf == null) {

            return DONE;

        } else {

            return cf;

        }

    }

    @Override

    public CompletableFuture<WebSocket> sendPing(ByteBuffer message) {

        Objects.requireNonNull(message);

        long id = 0;

        if (debug.on()) {

            id = sendCounter.incrementAndGet();

            debug.log("enter send ping %s payload=%s", id, message);

        }

        CompletableFuture<WebSocket> result;

        if (!setPendingPingOrPong()) {

            result = failedFuture(new IllegalStateException("Send pending"));

        } else {

            result = transport.sendPing(message, this,

                                        (r, e) -> clearPendingPingOrPong());

        }

        if (debug.on()) {

            debug.log("exit send ping %s returned %s", id, result);

        }

        return replaceNull(result);

    }

    @Override

    public CompletableFuture<WebSocket> sendPong(ByteBuffer message) {

        Objects.requireNonNull(message);

        long id = 0;

        if (debug.on()) {

            id = sendCounter.incrementAndGet();

            debug.log("enter send pong %s payload=%s", id, message);

        }

        CompletableFuture<WebSocket> result;

        if (!setPendingPingOrPong()) {

            result = failedFuture(new IllegalStateException("Send pending"));

        } else {

            result =  transport.sendPong(message, this,

                                         (r, e) -> clearPendingPingOrPong());

        }

        if (debug.on()) {

            debug.log("exit send pong %s returned %s", id, result);

        }

        return replaceNull(result);

    }

    private boolean setPendingPingOrPong() {

        return pendingPingOrPong.compareAndSet(false, true);

    }

    private void clearPendingPingOrPong() {

        pendingPingOrPong.set(false);

    }

    @Override

    public CompletableFuture<WebSocket> sendClose(int statusCode,

                                                  String reason) {

        Objects.requireNonNull(reason);

        long id = 0;

        if (debug.on()) {

            id = sendCounter.incrementAndGet();

            debug.log("enter send close %s statusCode=%s reason.length=%s",

                      id, statusCode, reason.length());

        }

        CompletableFuture<WebSocket> result;

        // Close message is the only type of message whose validity is checked

        // in the corresponding send method. This is made in order to close the

        // output in place. Otherwise the number of Close messages in queue

        // would not be bounded.

        if (!isLegalToSendFromClient(statusCode)) {

            result = failedFuture(new IllegalArgumentException("statusCode"));

        } else if (!isLegalReason(reason)) {

            result = failedFuture(new IllegalArgumentException("reason"));

        } else if (!outputClosed.compareAndSet(false, true)){

            result = failedFuture(new IOException("Output closed"));

        } else {

            result = sendClose0(statusCode, reason);

        }

        if (debug.on()) {

            debug.log("exit send close %s returned %s", id, result);

        }

        return replaceNull(result);

    }

    private static boolean isLegalReason(String reason) {

        if (reason.length() > 123) { // quick check

            return false;

        }

        CharsetEncoder encoder = StandardCharsets.UTF_8.newEncoder()

                        .onMalformedInput(CodingErrorAction.REPORT)

                        .onUnmappableCharacter(CodingErrorAction.REPORT);

        ByteBuffer bytes;

        try {

            bytes = encoder.encode(CharBuffer.wrap(reason));

        } catch (CharacterCodingException ignored) {

            return false;

        }

        return bytes.remaining() <= 123;

    }

    /*

     * The implementation uses this method internally to send Close messages

     * with codes that are not allowed to be sent through the API.

     */

    private CompletableFuture<WebSocket> sendClose0(int statusCode,

                                                    String reason) {

        return transport.sendClose(statusCode, reason, this,

                                   (r, e) -> processCloseError(e));

    }

    private void processCloseError(Throwable e) {

        if (e == null) {

            debug.log("send close completed successfully");

        } else {

            debug.log("send close completed with error", e);

        }

        outputClosed.set(true);

        try {

            transport.closeOutput();

        } catch (IOException ignored) { }

    }

    @Override

    public void request(long n) {

        if (debug.on()) {

            debug.log("request %s", n);

        }

        if (demand.increase(n)) {

            receiveScheduler.runOrSchedule();

        }

    }

    @Override

    public String getSubprotocol() {

        return subprotocol;

    }

    @Override

    public boolean isOutputClosed() {

        return outputClosed.get();

    }

    @Override

    public boolean isInputClosed() {

        return inputClosed;

    }

    @Override

    public void abort() {

        if (debug.on()) {

            debug.log("abort");

        }

        inputClosed = true;

        outputClosed.set(true);

        receiveScheduler.stop();

        close();

    }

    @Override

    public String toString() {

        return super.toString()

                + "[uri=" + uri

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

                + "]";

    }

    /*

     * The assumptions about order is as follows:

     *

     *     - state is never changed more than twice inside the `run` method:

     *       x --(1)--> IDLE --(2)--> y (otherwise we're loosing events, or

     *       overwriting parts of messages creating a mess since there's no

     *       queueing)

     *     - OPEN is always the first state

     *     - no messages are requested/delivered before onOpen is called (this

     *       is implemented by making WebSocket instance accessible first in

     *       onOpen)

     *     - after the state has been observed as CLOSE/ERROR, the scheduler

     *       is stopped

     */

    private class ReceiveTask extends SequentialScheduler.CompleteRestartableTask {

        // Transport only asked here and nowhere else because we must make sure

        // onOpen is invoked first and no messages become pending before onOpen

        // finishes

        @Override

        public void run() {

            if (debug.on()) {

                debug.log("enter receive task");

            }

            loop:

            while (!receiveScheduler.isStopped()) {

                State s = state.get();

                if (debug.on()) {

                    debug.log("receive state: %s", s);

                }

                try {

                    switch (s) {

                        case OPEN:

                            processOpen();

                            tryChangeState(OPEN, IDLE);

                            break;

                        case TEXT:

                            processText();

                            tryChangeState(TEXT, IDLE);

                            break;

                        case BINARY:

                            processBinary();

                            tryChangeState(BINARY, IDLE);

                            break;

                        case PING:

                            processPing();

                            tryChangeState(PING, IDLE);

                            break;

                        case PONG:

                            processPong();

                            tryChangeState(PONG, IDLE);

                            break;

                        case CLOSE:

                            processClose();

                            break loop;

                        case ERROR:

                            processError();

                            break loop;

                        case IDLE:

                            if (demand.tryDecrement()

                                    && tryChangeState(IDLE, WAITING)) {

                                transport.request(1);

                            }

                            break loop;

                        case WAITING:

                            // For debugging spurious signalling: when there was

                            // a signal, but apparently nothing has changed

                            break loop;

                        default:

                            throw new InternalError(String.valueOf(s));

                    }

                } catch (Throwable t) {

                    signalError(t);

                }

            }

            if (debug.on()) {

                debug.log("exit receive task");

            }

        }

        private void processError() throws IOException {

            if (debug.on()) {

                debug.log("processError");

            }

            transport.closeInput();

            receiveScheduler.stop();

            Throwable err = error.get();

            if (err instanceof FailWebSocketException) {

                int code1 = ((FailWebSocketException) err).getStatusCode();

                err = new ProtocolException().initCause(err);

                if (debug.on()) {

                    debug.log("failing %s with error=%s statusCode=%s",

                              WebSocketImpl.this, err, code1);

                }

                sendCloseSilently(code1);

            }

            long id = 0;

            if (debug.on()) {