/*
 * Copyright (c) 2015, 2016, 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 java.net.http;

import java.io.UncheckedIOException;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.Arrays;
import java.util.Base64;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;
import java.util.stream.Collectors;

import static java.lang.String.format;
import static java.lang.System.Logger.Level.TRACE;
import static java.net.http.WSUtils.logger;
import static java.net.http.WSUtils.webSocketSpecViolation;

final class WSOpeningHandshake {

    private static final String HEADER_CONNECTION = "Connection";
    private static final String HEADER_UPGRADE = "Upgrade";
    private static final String HEADER_ACCEPT = "Sec-WebSocket-Accept";
    private static final String HEADER_EXTENSIONS = "Sec-WebSocket-Extensions";
    private static final String HEADER_KEY = "Sec-WebSocket-Key";
    private static final String HEADER_PROTOCOL = "Sec-WebSocket-Protocol";
    private static final String HEADER_VERSION = "Sec-WebSocket-Version";
    private static final String VALUE_VERSION = "13"; // WebSocket's lucky number

    private static final SecureRandom srandom = new SecureRandom();

    private final MessageDigest sha1;

    {
        try {
            sha1 = MessageDigest.getInstance("SHA-1");
        } catch (NoSuchAlgorithmException e) {
            // Shouldn't happen:
            // SHA-1 must be available in every Java platform implementation
            throw new InternalError("Minimum platform requirements are not met", e);
        }
    }

    private final HttpRequest request;
    private final Collection<String> subprotocols;
    private final String nonce;

    WSOpeningHandshake(WSBuilder b) {
        URI httpURI = createHttpUri(b.getUri());
        HttpRequest.Builder requestBuilder = b.getClient().request(httpURI);
        Duration connectTimeout = b.getConnectTimeout();
        if (connectTimeout != null) {
            requestBuilder.timeout(
                    TimeUnit.of(ChronoUnit.MILLIS),
                    connectTimeout.get(ChronoUnit.MILLIS)
            );
        }
        Collection<String> s = b.getSubprotocols();
        if (!s.isEmpty()) {
            String p = s.stream().collect(Collectors.joining(", "));
            requestBuilder.header(HEADER_PROTOCOL, p);
        }
        requestBuilder.header(HEADER_VERSION, VALUE_VERSION);
        this.nonce = createNonce();
        requestBuilder.header(HEADER_KEY, this.nonce);
        this.request = requestBuilder.GET();
        HttpRequestImpl r = (HttpRequestImpl) this.request;
        r.isWebSocket(true);
        r.setSystemHeader(HEADER_UPGRADE, "websocket");
        r.setSystemHeader(HEADER_CONNECTION, "Upgrade");
        this.subprotocols = s;
    }

    private URI createHttpUri(URI webSocketUri) {
        // FIXME: check permission for WebSocket URI and translate it into http/https permission
        logger.log(TRACE, "->createHttpUri(''{0}'')", webSocketUri);
        String httpScheme = webSocketUri.getScheme().equalsIgnoreCase("ws")
                ? "http"
                : "https";
        try {
            URI uri = new URI(httpScheme,
                    webSocketUri.getUserInfo(),
                    webSocketUri.getHost(),
                    webSocketUri.getPort(),
                    webSocketUri.getPath(),
                    webSocketUri.getQuery(),
                    null);
            logger.log(TRACE, "<-createHttpUri: ''{0}''", uri);
            return uri;
        } catch (URISyntaxException e) {
            // Shouldn't happen: URI invariant
            throw new InternalError("Error translating WebSocket URI to HTTP URI", e);
        }
    }

    CompletableFuture<Result> performAsync() {
        // The whole dancing with thenCompose instead of thenApply is because
        // WebSocketHandshakeException is a checked exception
        return request.responseAsync()
                .thenCompose(response -> {
                    try {
                        Result result = handleResponse(response);
                        return CompletableFuture.completedFuture(result);
                    } catch (WebSocketHandshakeException e) {
                        return CompletableFuture.failedFuture(e);
                    } catch (UncheckedIOException ee) {
                        return CompletableFuture.failedFuture(ee.getCause());
                    }
                });
    }

    private Result handleResponse(HttpResponse response) throws WebSocketHandshakeException {
        // By this point all redirects, authentications, etc. (if any) must have
        // been done by the httpClient used by the WebSocket; so only 101 is
        // expected
        int statusCode = response.statusCode();
        if (statusCode != 101) {
            String m = webSocketSpecViolation("1.3.",
                    "Unable to complete handshake; HTTP response status code "
                            + statusCode
            );
            throw new WebSocketHandshakeException(m, response);
        }
        HttpHeaders h = response.headers();
        checkHeader(h, response, HEADER_UPGRADE, v -> v.equalsIgnoreCase("websocket"));
        checkHeader(h, response, HEADER_CONNECTION, v -> v.equalsIgnoreCase("Upgrade"));
        checkVersion(response, h);
        checkAccept(response, h);
        checkExtensions(response, h);
        String subprotocol = checkAndReturnSubprotocol(response, h);
        RawChannel channel = null;
        try {
            channel = ((HttpResponseImpl) response).rawChannel();
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
        return new Result(subprotocol, channel);
    }

    private void checkExtensions(HttpResponse response, HttpHeaders headers)
            throws WebSocketHandshakeException {
        List<String> ext = headers.allValues(HEADER_EXTENSIONS);
        if (!ext.isEmpty()) {
            String m = webSocketSpecViolation("4.1.",
                    "Server responded with extension(s) though none were requested "
                            + Arrays.toString(ext.toArray())
            );
            throw new WebSocketHandshakeException(m, response);
        }
    }

    private String checkAndReturnSubprotocol(HttpResponse response, HttpHeaders headers)
            throws WebSocketHandshakeException {
        assert response.statusCode() == 101 : response.statusCode();
        List<String> sp = headers.allValues(HEADER_PROTOCOL);
        int size = sp.size();
        if (size == 0) {
            // In this case the subprotocol requested (if any) by the client
            // doesn't matter. If there is no such header in the response, then
            // the server doesn't want to use any subprotocol
            return null;
        } else if (size > 1) {
            // We don't know anything about toString implementation of this
            // list, so let's create an array
            String m = webSocketSpecViolation("4.1.",
                    "Server responded with multiple subprotocols: "
                            + Arrays.toString(sp.toArray())
            );
            throw new WebSocketHandshakeException(m, response);
        } else {
            String selectedSubprotocol = sp.get(0);
            if (this.subprotocols.contains(selectedSubprotocol)) {
                return selectedSubprotocol;
            } else {
                String m = webSocketSpecViolation("4.1.",
                        format("Server responded with a subprotocol " +
                                        "not among those requested: '%s'",
                                selectedSubprotocol));
                throw new WebSocketHandshakeException(m, response);
            }
        }
    }

    private void checkAccept(HttpResponse response, HttpHeaders headers)
            throws WebSocketHandshakeException {
        assert response.statusCode() == 101 : response.statusCode();
        String x = nonce + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
        sha1.update(x.getBytes(StandardCharsets.ISO_8859_1));
        String expected = Base64.getEncoder().encodeToString(sha1.digest());
        checkHeader(headers, response, HEADER_ACCEPT, actual -> actual.trim().equals(expected));
    }

    private void checkVersion(HttpResponse response, HttpHeaders headers)
            throws WebSocketHandshakeException {
        assert response.statusCode() == 101 : response.statusCode();
        List<String> versions = headers.allValues(HEADER_VERSION);
        if (versions.isEmpty()) { // That's normal and expected
            return;
        }
        String m = webSocketSpecViolation("4.4.",
                "Server responded with version(s) "
                        + Arrays.toString(versions.toArray()));
        throw new WebSocketHandshakeException(m, response);
    }

    //
    // Checks whether there's only one value for the header with the given name
    // and the value satisfies the predicate.
    //
    private static void checkHeader(HttpHeaders headers,
                                    HttpResponse response,
                                    String headerName,
                                    Predicate<? super String> valuePredicate)
            throws WebSocketHandshakeException {
        assert response.statusCode() == 101 : response.statusCode();
        List<String> values = headers.allValues(headerName);
        if (values.isEmpty()) {
            String m = webSocketSpecViolation("4.1.",
                    format("Server response field '%s' is missing", headerName)
            );
            throw new WebSocketHandshakeException(m, response);
        } else if (values.size() > 1) {
            String m = webSocketSpecViolation("4.1.",
                    format("Server response field '%s' has multiple values", headerName)
            );
            throw new WebSocketHandshakeException(m, response);
        }
        if (!valuePredicate.test(values.get(0))) {
            String m = webSocketSpecViolation("4.1.",
                    format("Server response field '%s' is incorrect", headerName)
            );
            throw new WebSocketHandshakeException(m, response);
        }
    }

    private static String createNonce() {
        byte[] bytes = new byte[16];
        srandom.nextBytes(bytes);
        return Base64.getEncoder().encodeToString(bytes);
    }

    static final class Result {

        final String subprotocol;
        final RawChannel channel;

        private Result(String subprotocol, RawChannel channel) {
            this.subprotocol = subprotocol;
            this.channel = channel;
        }
    }
}