/*
 * ====================================================================
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 * ====================================================================
 *
 * This software consists of voluntary contributions made by many
 * individuals on behalf of the Apache Software Foundation.  For more
 * information on the Apache Software Foundation, please see
 * <http://www.apache.org/>.
 *
 */
package org.apache.hc.client5.http.websocket.client.impl.protocol;

import java.io.IOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.SecureRandom;
import java.util.Base64;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.StringJoiner;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;

import org.apache.hc.client5.http.classic.methods.HttpGet;
import org.apache.hc.client5.http.websocket.api.WebSocket;
import org.apache.hc.client5.http.websocket.api.WebSocketClientConfig;
import org.apache.hc.client5.http.websocket.api.WebSocketListener;
import org.apache.hc.client5.http.websocket.client.impl.connector.WebSocketEndpointConnector;
import org.apache.hc.core5.websocket.extension.ExtensionChain;
import org.apache.hc.core5.websocket.extension.PerMessageDeflate;
import org.apache.hc.client5.http.websocket.transport.WebSocketUpgrader;
import org.apache.hc.core5.annotation.Internal;
import org.apache.hc.core5.concurrent.FutureCallback;
import org.apache.hc.core5.http.EntityDetails;
import org.apache.hc.core5.http.Header;
import org.apache.hc.core5.http.HttpException;
import org.apache.hc.core5.http.HttpHeaders;
import org.apache.hc.core5.http.HttpHost;
import org.apache.hc.core5.http.HttpResponse;
import org.apache.hc.core5.http.HttpStatus;
import org.apache.hc.core5.http.URIScheme;
import org.apache.hc.core5.http.impl.bootstrap.HttpAsyncRequester;
import org.apache.hc.core5.http.message.BasicHttpRequest;
import org.apache.hc.core5.http.nio.AsyncClientExchangeHandler;
import org.apache.hc.core5.http.nio.CapacityChannel;
import org.apache.hc.core5.http.nio.DataStreamChannel;
import org.apache.hc.core5.http.nio.RequestChannel;
import org.apache.hc.core5.http.protocol.HttpContext;
import org.apache.hc.core5.pool.ManagedConnPool;
import org.apache.hc.core5.reactor.IOSession;
import org.apache.hc.core5.reactor.ProtocolIOSession;
import org.apache.hc.core5.util.Args;
import org.apache.hc.core5.util.Timeout;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * HTTP/1.1 Upgrade (RFC 6455). Uses getters on WebSocketClientConfig.
 */
@Internal
public final class Http1UpgradeProtocol implements WebSocketProtocolStrategy {

    private static final Logger LOG = LoggerFactory.getLogger(Http1UpgradeProtocol.class);
    private static final SecureRandom NONCE_RANDOM = new SecureRandom();

    private final HttpAsyncRequester requester;
    private final ManagedConnPool<HttpHost, IOSession> connPool;

    public Http1UpgradeProtocol(final HttpAsyncRequester requester, final ManagedConnPool<HttpHost, IOSession> connPool) {
        this.requester = requester;
        this.connPool = connPool;
    }

    @Override
    public CompletableFuture<WebSocket> connect(
            final URI uri,
            final WebSocketListener listener,
            final WebSocketClientConfig cfg,
            final HttpContext context) {

        Args.notNull(uri, "uri");
        Args.notNull(listener, "listener");
        Args.notNull(cfg, "cfg");

        final boolean secure = "wss".equalsIgnoreCase(uri.getScheme());
        if (!secure && !"ws".equalsIgnoreCase(uri.getScheme())) {
            final CompletableFuture<WebSocket> f = new CompletableFuture<>();
            f.completeExceptionally(new IllegalArgumentException("Scheme must be ws or wss"));
            return f;
        }

        final String scheme = secure ? URIScheme.HTTPS.id : URIScheme.HTTP.id;
        final int port = uri.getPort() > 0 ? uri.getPort() : secure ? 443 : 80;
        final String host = Args.notBlank(uri.getHost(), "host");
        String path = uri.getRawPath();
        if (path == null || path.isEmpty()) {
            path = "/";
        }
        final String fullPath = uri.getRawQuery() != null ? path + "?" + uri.getRawQuery() : path;
        final HttpHost target = new HttpHost(scheme, host, port);

        final CompletableFuture<WebSocket> result = new CompletableFuture<>();
        final WebSocketEndpointConnector wsRequester = new WebSocketEndpointConnector(requester, connPool);

        final Timeout timeout = cfg.getConnectTimeout() != null ? cfg.getConnectTimeout() : Timeout.ofSeconds(10);

        wsRequester.connect(target, timeout, null,
                new FutureCallback<WebSocketEndpointConnector.ProtoEndpoint>() {
                    @Override
                    public void completed(final WebSocketEndpointConnector.ProtoEndpoint endpoint) {
                        try {
                            final String secKey = randomKey();
                            final BasicHttpRequest req = new BasicHttpRequest(HttpGet.METHOD_NAME, target, fullPath);

                            req.addHeader(HttpHeaders.CONNECTION, "Upgrade");
                            req.addHeader(HttpHeaders.UPGRADE, "websocket");
                            req.addHeader("Sec-WebSocket-Version", "13");
                            req.addHeader("Sec-WebSocket-Key", secKey);

                            // subprotocols
                            if (!cfg.getSubprotocols().isEmpty()) {
                                final StringJoiner sj = new StringJoiner(", ");
                                for (final String p : cfg.getSubprotocols()) {
                                    if (p != null && !p.isEmpty()) {
                                        sj.add(p);
                                    }
                                }
                                final String offered = sj.toString();
                                if (!offered.isEmpty()) {
                                    req.addHeader("Sec-WebSocket-Protocol", offered);
                                }
                            }

                            // PMCE offer
                            if (cfg.isPerMessageDeflateEnabled()) {
                                final StringBuilder ext = new StringBuilder("permessage-deflate");
                                if (cfg.isOfferServerNoContextTakeover()) {
                                    ext.append("; server_no_context_takeover");
                                }
                                if (cfg.isOfferClientNoContextTakeover()) {
                                    ext.append("; client_no_context_takeover");
                                }
                                // client_max_window_bits is never offered: the JDK Deflater cannot
                                // honor a server-selected window smaller than 15 (RFC 7692 7.2.1).
                                if (cfg.getOfferServerMaxWindowBits() != null) {
                                    ext.append("; server_max_window_bits=").append(cfg.getOfferServerMaxWindowBits());
                                }
                                req.addHeader("Sec-WebSocket-Extensions", ext.toString());
                            }

                            if (LOG.isDebugEnabled()) {
                                LOG.debug("Dispatching HTTP/1.1 Upgrade: GET {} with headers:", fullPath);
                                for (final Header h : req.getHeaders()) {
                                    LOG.debug("  {}: {}", h.getName(), h.getValue());
                                }
                            }

                            final AtomicBoolean done = new AtomicBoolean(false);

                            final AsyncClientExchangeHandler upgrade = new AsyncClientExchangeHandler() {
                                @Override
                                public void releaseResources() {
                                }

                                @Override
                                public void failed(final Exception cause) {
                                    if (done.compareAndSet(false, true)) {
                                        try {
                                            endpoint.releaseAndDiscard();
                                        } catch (final Throwable ignore) {
                                        }
                                        result.completeExceptionally(cause);
                                    }
                                }

                                @Override
                                public void cancel() {
                                    if (done.compareAndSet(false, true)) {
                                        try {
                                            endpoint.releaseAndDiscard();
                                        } catch (final Throwable ignore) {
                                        }
                                        result.cancel(true);
                                    }
                                }

                                @Override
                                public void produceRequest(final RequestChannel ch,
                                                           final HttpContext hc)
                                        throws IOException, HttpException {
                                    ch.sendRequest(req, null, hc);
                                }

                                @Override
                                public int available() {
                                    return 0;
                                }

                                @Override
                                public void produce(final DataStreamChannel channel) {
                                }

                                @Override
                                public void updateCapacity(final CapacityChannel capacityChannel) {
                                }

                                @Override
                                public void consume(final ByteBuffer src) {
                                }

                                @Override
                                public void streamEnd(final List<? extends Header> trailers) {
                                }

                                @Override
                                public void consumeInformation(final HttpResponse response,
                                                               final HttpContext hc) {
                                    final int code = response.getCode();
                                    if (code == HttpStatus.SC_SWITCHING_PROTOCOLS && done.compareAndSet(false, true)) {
                                        finishUpgrade(endpoint, response, secKey, listener, cfg, result);
                                    }
                                }

                                @Override
                                public void consumeResponse(final HttpResponse response,
                                                            final EntityDetails entity,
                                                            final HttpContext hc) {
                                    final int code = response.getCode();
                                    if (code == HttpStatus.SC_SWITCHING_PROTOCOLS && done.compareAndSet(false, true)) {
                                        finishUpgrade(endpoint, response, secKey, listener, cfg, result);
                                        return;
                                    }
                                    failed(new IllegalStateException("Unexpected status: " + code));
                                }
                            };

                            endpoint.execute(upgrade, null, context);

                        } catch (final Exception ex) {
                            try {
                                endpoint.releaseAndDiscard();
                            } catch (final Throwable ignore) {
                            }
                            result.completeExceptionally(ex);
                        }
                    }

                    @Override
                    public void failed(final Exception ex) {
                        result.completeExceptionally(ex);
                    }

                    @Override
                    public void cancelled() {
                        result.cancel(true);
                    }
                });

        return result;
    }

    private void finishUpgrade(
            final WebSocketEndpointConnector.ProtoEndpoint endpoint,
            final HttpResponse response,
            final String secKey,
            final WebSocketListener listener,
            final WebSocketClientConfig cfg,
            final CompletableFuture<WebSocket> result) {
        try {
            final String accept = headerValue(response, "Sec-WebSocket-Accept");
            final String expected = expectedAccept(secKey);
            final String acceptValue = accept != null ? accept.trim() : null;
            if (!expected.equals(acceptValue)) {
                throw new IllegalStateException("Bad Sec-WebSocket-Accept");
            }

            final String upgrade = headerValue(response, "Upgrade");
            if (upgrade == null || !"websocket".equalsIgnoreCase(upgrade.trim())) {
                throw new IllegalStateException("Missing/invalid Upgrade header: " + upgrade);
            }
            if (!containsToken(response, "Connection", "Upgrade")) {
                throw new IllegalStateException("Missing/invalid Connection header");
            }

            final String proto = headerValue(response, "Sec-WebSocket-Protocol");
            if (proto != null && !proto.isEmpty()) {
                boolean matched = false;
                for (final String p : cfg.getSubprotocols()) {
                    if (p.equals(proto)) {
                        matched = true;
                        break;
                    }
                }
                if (!matched) {
                    throw new IllegalStateException("Server selected subprotocol not offered: " + proto);
                }
            }

            final ExtensionChain chain = buildExtensionChain(cfg, headerValue(response, "Sec-WebSocket-Extensions"));

            final ProtocolIOSession ioSession = endpoint.getProtocolIOSession();
            final WebSocketUpgrader upgrader = new WebSocketUpgrader(listener, cfg, chain, endpoint);
            ioSession.registerProtocol("websocket", upgrader);
            ioSession.switchProtocol("websocket", new FutureCallback<ProtocolIOSession>() {
                @Override
                public void completed(final ProtocolIOSession s) {
                    s.setSocketTimeout(Timeout.DISABLED);
                    final WebSocket ws = upgrader.getWebSocket();
                    try {
                        listener.onOpen(ws);
                    } catch (final Throwable ignore) {
                    }
                    result.complete(ws);
                }

                @Override
                public void failed(final Exception ex) {
                    try {
                        endpoint.releaseAndDiscard();
                    } catch (final Throwable ignore) {
                    }
                    result.completeExceptionally(ex);
                }

                @Override
                public void cancelled() {
                    try {
                        endpoint.releaseAndDiscard();
                    } catch (final Throwable ignore) {
                    }
                    result.cancel(true);
                }
            });

        } catch (final Exception ex) {
            try {
                endpoint.releaseAndDiscard();
            } catch (final Throwable ignore) {
            }
            result.completeExceptionally(ex);
        }
    }

    static String headerValue(final HttpResponse r, final String name) {
        // RFC 6455 section 9.1 permits Sec-WebSocket-Extensions and Sec-WebSocket-Protocol
        // to be split across multiple header fields; combine them as a comma-separated list.
        final Header[] headers = r.getHeaders(name);
        if (headers.length == 0) {
            return null;
        }
        if (headers.length == 1) {
            return headers[0].getValue();
        }
        final StringBuilder buf = new StringBuilder();
        for (final Header h : headers) {
            final String value = h.getValue();
            if (value == null || value.isEmpty()) {
                continue;
            }
            if (buf.length() > 0) {
                buf.append(", ");
            }
            buf.append(value);
        }
        return buf.length() > 0 ? buf.toString() : null;
    }

    private static boolean containsToken(final HttpResponse r, final String header, final String token) {
        for (final Header h : r.getHeaders(header)) {
            for (final String p : h.getValue().split(",")) {
                if (p.trim().equalsIgnoreCase(token)) {
                    return true;
                }
            }
        }
        return false;
    }

    private static String randomKey() {
        final byte[] nonce = new byte[16];
        NONCE_RANDOM.nextBytes(nonce);
        return Base64.getEncoder().encodeToString(nonce);
    }

    private static String expectedAccept(final String key) throws Exception {
        final MessageDigest sha1 = MessageDigest.getInstance("SHA-1");
        sha1.update((key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").getBytes(StandardCharsets.US_ASCII));
        return Base64.getEncoder().encodeToString(sha1.digest());
    }

    static ExtensionChain buildExtensionChain(final WebSocketClientConfig cfg, final String ext) {
        final ExtensionChain chain = new ExtensionChain();
        if (ext == null || ext.isEmpty()) {
            return chain;
        }
        boolean pmceSeen = false, serverNoCtx = false, clientNoCtx = false;
        Integer serverBits = null;
        final Integer offerServerBits = cfg.getOfferServerMaxWindowBits();

        final String[] tokens = ext.split(",");
        for (final String raw0 : tokens) {
            final String raw = raw0.trim();
            final String[] parts = raw.split(";");
            final String token = parts[0].trim().toLowerCase();

            // Only permessage-deflate is supported
            if (!"permessage-deflate".equals(token)) {
                throw new IllegalStateException("Server selected unsupported extension: " + token);
            }
            if (pmceSeen) {
                throw new IllegalStateException("Server selected permessage-deflate more than once");
            }
            pmceSeen = true;

            // RFC 7692 section 7.1: a parameter name must not appear more than once within an
            // accepted extension; a repeated name makes the response invalid.
            final Set<String> parameterNames = new HashSet<>();
            for (int i = 1; i < parts.length; i++) {
                final String p = parts[i].trim();
                final int eq = p.indexOf('=');
                final String parameterName = (eq >= 0 ? p.substring(0, eq) : p).trim().toLowerCase(Locale.ROOT);
                if (!parameterNames.add(parameterName)) {
                    throw new IllegalStateException("Duplicate permessage-deflate parameter: " + parameterName);
                }
                if (eq < 0) {
                    // RFC 7692 sections 7.1.1.1 and 7.1.1.2 permit the server to include either
                    // no-context-takeover parameter in the response even when the offer did not;
                    // both are always safe to honour.
                    if ("server_no_context_takeover".equalsIgnoreCase(p)) {
                        serverNoCtx = true;
                    } else if ("client_no_context_takeover".equalsIgnoreCase(p)) {
                        clientNoCtx = true;
                    } else {
                        throw new IllegalStateException("Unsupported permessage-deflate parameter: " + p);
                    }
                } else {
                    final String k = p.substring(0, eq).trim();
                    String v = p.substring(eq + 1).trim();
                    if (v.length() >= 2 && v.charAt(0) == '"' && v.charAt(v.length() - 1) == '"') {
                        v = v.substring(1, v.length() - 1); // strip quotes if any
                    }
                    if ("client_max_window_bits".equalsIgnoreCase(k)) {
                        // The client never offers client_max_window_bits (the JDK Deflater cannot honor
                        // a window smaller than 15, RFC 7692 section 7.2.1), so a server that selects it
                        // has echoed a parameter that was never offered.
                        throw new IllegalStateException("Server selected client_max_window_bits not offered");
                    } else if ("server_max_window_bits".equalsIgnoreCase(k)) {
                        // RFC 7692 section 7.1.2.1 permits the server to include this parameter
                        // uninvited; a server constraining its own window is always acceptable.
                        try {
                            if (v.isEmpty()) {
                                throw new IllegalStateException("server_max_window_bits must have a value");
                            }
                            serverBits = Integer.parseInt(v);
                            if (serverBits < 8 || serverBits > 15) {
                                throw new IllegalStateException("server_max_window_bits out of range: " + serverBits);
                            }
                        } catch (final NumberFormatException nfe) {
                            throw new IllegalStateException("Invalid server_max_window_bits: " + v, nfe);
                        }
                    } else {
                        throw new IllegalStateException("Unsupported permessage-deflate parameter: " + k);
                    }
                }
            }
        }

        if (pmceSeen) {
            if (!cfg.isPerMessageDeflateEnabled()) {
                throw new IllegalStateException("Server negotiated PMCE but client disabled it");
            }
            if (serverBits != null) {
                if (offerServerBits != null && serverBits > offerServerBits) {
                    throw new IllegalStateException("server_max_window_bits exceeds offer: " + serverBits);
                }
            }
            chain.add(new PerMessageDeflate(true, serverNoCtx, clientNoCtx, null, serverBits));
        }
        return chain;
    }
}
