// Licensed to the Software Freedom Conservancy (SFC) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The SFC 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.

package org.openqa.selenium.remote.http.jdk;

import com.google.auto.service.AutoService;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.net.Authenticator;
import java.net.PasswordAuthentication;
import java.net.ProtocolException;
import java.net.Proxy;
import java.net.ProxySelector;
import java.net.SocketAddress;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.http.HttpClient.Version;
import java.net.http.HttpResponse.BodyHandler;
import java.net.http.HttpResponse.BodyHandlers;
import java.net.http.HttpTimeoutException;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Objects;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.net.ssl.SSLContext;
import org.openqa.selenium.Credentials;
import org.openqa.selenium.TimeoutException;
import org.openqa.selenium.UsernameAndPassword;
import org.openqa.selenium.WebDriverException;
import org.openqa.selenium.remote.http.BinaryMessage;
import org.openqa.selenium.remote.http.ClientConfig;
import org.openqa.selenium.remote.http.CloseMessage;
import org.openqa.selenium.remote.http.ConnectionFailedException;
import org.openqa.selenium.remote.http.HttpClient;
import org.openqa.selenium.remote.http.HttpClientName;
import org.openqa.selenium.remote.http.HttpHandler;
import org.openqa.selenium.remote.http.HttpMethod;
import org.openqa.selenium.remote.http.HttpRequest;
import org.openqa.selenium.remote.http.HttpResponse;
import org.openqa.selenium.remote.http.Message;
import org.openqa.selenium.remote.http.TextMessage;
import org.openqa.selenium.remote.http.WebSocket;

public class JdkHttpClient implements HttpClient {
  public static final Logger LOG = Logger.getLogger(JdkHttpClient.class.getName());
  private static final AtomicInteger POOL_COUNTER = new AtomicInteger(0);
  private final JdkHttpMessages messages;
  private final HttpHandler handler;
  private java.net.http.HttpClient client;
  private final List<WebSocket> websockets;
  private final ExecutorService executorService;
  private final Duration readTimeout;
  private final Duration connectTimeout;

  JdkHttpClient(ClientConfig config) {
    Objects.requireNonNull(config, "Client config must be set");

    this.messages = new JdkHttpMessages(config);
    this.readTimeout = config.readTimeout();
    this.connectTimeout = config.connectionTimeout();
    this.websockets = new ArrayList<>();
    this.handler = config.filter().andFinally(this::execute0);

    String poolName = "JdkHttpClient-" + POOL_COUNTER.getAndIncrement();
    AtomicInteger threadCounter = new AtomicInteger(0);
    executorService =
        Executors.newCachedThreadPool(
            r -> {
              Thread thread = new Thread(r, poolName + "-" + threadCounter.getAndIncrement());
              thread.setDaemon(true);
              return thread;
            });

    java.net.http.HttpClient.Builder builder =
        java.net.http.HttpClient.newBuilder()
            .connectTimeout(connectTimeout)
            .followRedirects(java.net.http.HttpClient.Redirect.NEVER)
            .executor(executorService);

    Credentials credentials = config.credentials();
    String info = config.baseUri().getUserInfo();
    if (info != null && !info.trim().isEmpty()) {
      String[] parts = info.split(":", 2);
      String username = parts[0];
      String password = parts.length > 1 ? parts[1] : null;

      Authenticator authenticator =
          new Authenticator() {
            @Override
            protected PasswordAuthentication getPasswordAuthentication() {
              return new PasswordAuthentication(username, password.toCharArray());
            }
          };
      builder = builder.authenticator(authenticator);
    } else if (credentials != null) {
      if (!(credentials instanceof UsernameAndPassword)) {
        throw new IllegalArgumentException(
            "Credentials must be a user name and password: " + credentials);
      }
      UsernameAndPassword uap = (UsernameAndPassword) credentials;
      Authenticator authenticator =
          new Authenticator() {
            @Override
            protected PasswordAuthentication getPasswordAuthentication() {
              return new PasswordAuthentication(uap.username(), uap.password().toCharArray());
            }
          };
      builder = builder.authenticator(authenticator);
    }

    Proxy proxy = config.proxy();
    if (proxy != null) {
      ProxySelector proxySelector = new HttpProxySelector(proxy);
      builder = builder.proxy(proxySelector);
    }

    SSLContext sslContext = config.sslContext();
    if (sslContext != null) {
      builder.sslContext(sslContext);
    }

    String version = config.version();
    if (version != null) {
      builder.version(Version.valueOf(version));
    }

    this.client = builder.build();
  }

  @Override
  public WebSocket openSocket(HttpRequest request, WebSocket.Listener listener) {
    URI uri;

    try {
      uri = getWebSocketUri(request);
    } catch (URISyntaxException e) {
      throw new ConnectionFailedException("JdkWebSocket initial request execution error", e);
    }

    java.net.http.WebSocket.Builder builder = client.newWebSocketBuilder();

    request.getHeaderNames().forEach(name -> builder.header(name, request.getHeader(name)));

    CompletableFuture<Integer> closed = new CompletableFuture<>();
    CompletableFuture<java.net.http.WebSocket> webSocketCompletableFuture =
        builder
            .connectTimeout(connectTimeout)
            .buildAsync(
                uri,
                new java.net.http.WebSocket.Listener() {
                  final StringBuilder builder = new StringBuilder();
                  final ByteArrayOutputStream buffer = new ByteArrayOutputStream();

                  @Override
                  public CompletionStage<?> onText(
                      java.net.http.WebSocket webSocket, CharSequence data, boolean last) {
                    LOG.fine("Text message received. Appending data");
                    builder.append(data);

                    if (last) {
                      LOG.log(
                          Level.FINE,
                          "Final part of text message received. Calling listener with {0}",
                          builder);
                      listener.onText(builder.toString());
                      builder.setLength(0);
                    }

                    webSocket.request(1);
                    return null;
                  }

                  @Override
                  public CompletionStage<?> onBinary(
                      java.net.http.WebSocket webSocket, ByteBuffer data, boolean last) {
                    LOG.fine("Binary data received. Appending data");
                    byte[] ary = new byte[8192];

                    while (data.hasRemaining()) {
                      int n = Math.min(ary.length, data.remaining());
                      data.get(ary, 0, n);
                      buffer.write(ary, 0, n);
                    }

                    if (last) {
                      LOG.log(
                          Level.FINE,
                          "Final part of binary data received. Calling listener with {0} bytes of"
                              + " data",
                          buffer.size());
                      listener.onBinary(buffer.toByteArray());
                      buffer.reset();
                    }
                    webSocket.request(1);
                    return null;
                  }

                  @Override
                  public CompletionStage<?> onClose(
                      java.net.http.WebSocket webSocket, int statusCode, String reason) {
                    LOG.fine("Closing websocket");
                    closed.complete(statusCode);
                    listener.onClose(statusCode, reason);
                    return null;
                  }

                  @Override
                  public void onError(java.net.http.WebSocket webSocket, Throwable error) {
                    LOG.log(Level.FINE, "An error has occurred: " + error.getMessage(), error);
                    listener.onError(error);
                  }
                });

    java.net.http.WebSocket underlyingSocket;

    try {
      underlyingSocket =
          webSocketCompletableFuture.get(readTimeout.toMillis(), TimeUnit.MILLISECONDS);
    } catch (CancellationException e) {
      throw new ConnectionFailedException("JdkWebSocket initial request canceled", e);
    } catch (ExecutionException e) {
      Throwable cause = e.getCause();
      throw new ConnectionFailedException(
          "JdkWebSocket initial request execution error", (cause != null) ? cause : e);
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      throw new ConnectionFailedException("JdkWebSocket initial request interrupted", e);
    } catch (java.util.concurrent.TimeoutException e) {
      webSocketCompletableFuture.cancel(true);
      throw new ConnectionFailedException("JdkWebSocket initial request timeout", e);
    }

    WebSocket websocket =
        new WebSocket() {
          @Override
          public WebSocket send(Message message) {
            Supplier<CompletableFuture<java.net.http.WebSocket>> makeCall;

            if (message instanceof BinaryMessage) {
              BinaryMessage binaryMessage = (BinaryMessage) message;
              LOG.fine("Sending binary message");
              makeCall =
                  () -> underlyingSocket.sendBinary(ByteBuffer.wrap(binaryMessage.data()), true);
            } else if (message instanceof TextMessage) {
              TextMessage textMessage = (TextMessage) message;
              LOG.log(Level.FINE, "Sending text message: {0}", textMessage.text());
              makeCall = () -> underlyingSocket.sendText(textMessage.text(), true);
            } else if (message instanceof CloseMessage) {
              CloseMessage closeMessage = (CloseMessage) message;
              if (!underlyingSocket.isOutputClosed()) {
                // Sometimes the statusCode is -1, which could mean the socket is already closed.
                // We send a normal closure code in that case, though
                int statusCode = closeMessage.code() == -1 ? 1000 : closeMessage.code();
                LOG.log(
                    Level.FINE,
                    "Sending close message, statusCode {0}, reason: {1}",
                    new Object[] {statusCode, closeMessage.reason()});
                makeCall =
                    () -> {
                      CompletableFuture<java.net.http.WebSocket> future =
                          underlyingSocket.sendClose(statusCode, closeMessage.reason());
                      try {
                        closed.get(4, TimeUnit.SECONDS);
                      } catch (ExecutionException e) {
                        LOG.log(
                            Level.WARNING,
                            "failed to wait for the websocket to close",
                            e.getCause());
                      } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                      } catch (java.util.concurrent.TimeoutException e) {
                        LOG.finer("wait for the websocket to close timed out");
                      }
                      return future;
                    };
              } else {
                LOG.fine("Output is closed, not sending close message");
                return this;
              }
            } else {
              throw new IllegalArgumentException("Unsupported message type: " + message);
            }

            synchronized (underlyingSocket) {
              long start = System.currentTimeMillis();
              CompletableFuture<java.net.http.WebSocket> future = makeCall.get();
              try {
                future.get(readTimeout.toMillis(), TimeUnit.MILLISECONDS);
              } catch (CancellationException e) {
                throw new WebDriverException(e.getMessage(), e);
              } catch (ExecutionException e) {
                Throwable cause = e.getCause();
                if (cause == null) {
                  throw new WebDriverException(e);
                }
                throw new WebDriverException(cause);
              } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new WebDriverException(e.getMessage());
              } catch (java.util.concurrent.TimeoutException e) {
                throw new TimeoutException(e);
              } finally {
                LOG.log(
                    Level.FINE,
                    "Websocket response to {0} read in {1}ms",
                    new Object[] {message, (System.currentTimeMillis() - start)});
              }
            }
            return this;
          }

          @Override
          public void close() {
            LOG.fine("Closing websocket");
            send(new CloseMessage(1000, "WebDriver closing socket"));
          }
        };
    websockets.add(websocket);
    return websocket;
  }

  private URI getWebSocketUri(HttpRequest request) throws URISyntaxException {
    URI uri = messages.getRawUri(request);
    if ("http".equalsIgnoreCase(uri.getScheme())) {
      uri =
          new URI(
              "ws",
              uri.getUserInfo(),
              uri.getHost(),
              uri.getPort(),
              uri.getPath(),
              uri.getQuery(),
              uri.getFragment());
    } else if ("https".equalsIgnoreCase(uri.getScheme())) {
      uri =
          new URI(
              "wss",
              uri.getUserInfo(),
              uri.getHost(),
              uri.getPort(),
              uri.getPath(),
              uri.getQuery(),
              uri.getFragment());
    }
    return uri;
  }

  @Override
  public CompletableFuture<HttpResponse> executeAsync(HttpRequest request) {
    // the facade for this http request
    CompletableFuture<HttpResponse> cf = new CompletableFuture<>();

    // the actual http request
    Future<?> future =
        executorService.submit(
            () -> {
              try {
                HttpResponse response = handler.execute(request);

                cf.complete(response);
              } catch (Throwable t) {
                cf.completeExceptionally(t);
              }
            });

    cf.whenComplete(
        (result, throwable) -> {
          if (throwable instanceof java.util.concurrent.CancellationException) {
            // try to interrupt the http request in case someone canceled the future returned
            future.cancel(true);
          } else if (throwable instanceof java.util.concurrent.TimeoutException) {
            // try to interrupt the http request in case of a timeout, to avoid
            // https://bugs.openjdk.org/browse/JDK-8258397
            future.cancel(true);
          }
        });

    // will complete exceptionally with a java.util.concurrent.TimeoutException
    return cf.orTimeout(readTimeout.toMillis(), TimeUnit.MILLISECONDS);
  }

  @Override
  public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
    Future<HttpResponse> async = executeAsync(req);
    try {
      // executeAsync does define a timeout, no need to use a timeout here
      return async.get();
    } catch (CancellationException e) {
      throw new WebDriverException(e.getMessage(), e);
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      async.cancel(true);
      throw new WebDriverException(e.getMessage(), e);
    } catch (ExecutionException e) {
      Throwable cause = e.getCause();

      if (cause instanceof java.util.concurrent.TimeoutException) {
        throw new TimeoutException(cause);
      } else if (cause instanceof RuntimeException) {
        throw (RuntimeException) cause;
      } else if (cause instanceof Error) {
        throw (Error) cause;
      }

      throw new WebDriverException((cause != null) ? cause : e);
    }
  }

  private HttpResponse execute0(HttpRequest req) throws UncheckedIOException {
    Objects.requireNonNull(req, "Request");

    LOG.log(Level.FINE, "Executing request: {0}", req);
    long start = System.currentTimeMillis();

    BodyHandler<byte[]> byteHandler = BodyHandlers.ofByteArray();
    try {
      HttpMethod method = req.getMethod();
      URI rawUri = messages.getRawUri(req);

      // We need a custom handling of redirects to:
      // - increase the maximum number of retries to 100
      // - avoid a downgrade of POST requests, see the javadoc of j.n.h.HttpClient.Redirect
      // - not run into https://bugs.openjdk.org/browse/JDK-8304701
      for (int i = 0; i < 100; i++) {
        if (Thread.interrupted()) {
          throw new InterruptedException("http request has been interrupted");
        }

        java.net.http.HttpRequest request = messages.createRequest(req, method, rawUri);
        java.net.http.HttpResponse<byte[]> response = client.send(request, byteHandler);

        switch (response.statusCode()) {
          case 303:
            method = HttpMethod.GET;
          // fall-through
          case 301:
          case 302:
          case 307:
          case 308:
            URI location =
                rawUri.resolve(
                    response
                        .headers()
                        .firstValue("location")
                        .orElseThrow(
                            () ->
                                new ProtocolException(
                                    "HTTP "
                                        + response.statusCode()
                                        + " without 'location' header set")));

            if ("https".equalsIgnoreCase(rawUri.getScheme())
                && !"https".equalsIgnoreCase(location.getScheme())) {
              throw new SecurityException("Downgrade from secure to insecure connection.");
            } else if ("wss".equalsIgnoreCase(rawUri.getScheme())
                && !"wss".equalsIgnoreCase(location.getScheme())) {
              throw new SecurityException("Downgrade from secure to insecure connection.");
            }

            rawUri = location;
            continue;
          default:
            return messages.createResponse(response);
        }
      }

      throw new ProtocolException("Too many redirects: 101");
    } catch (HttpTimeoutException e) {
      throw new TimeoutException(e);
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      throw new WebDriverException(e.getMessage(), e);
    } finally {
      LOG.log(
          Level.FINE,
          "Ending request {0} in {1}ms",
          new Object[] {req, (System.currentTimeMillis() - start)});
    }
  }

  @Override
  public void close() {
    if (this.client == null) {
      return;
    }

    for (WebSocket websocket : websockets) {
      try {
        websocket.close();
      } catch (Exception e) {
        LOG.log(Level.WARNING, "failed to close the websocket: " + websocket, e);
      }
    }

    if (this.client instanceof AutoCloseable) {
      AutoCloseable closeable = (AutoCloseable) this.client;
      executorService.submit(
          () -> {
            try {
              closeable.close();
            } catch (Exception e) {
              LOG.log(Level.WARNING, "failed to close the http client: " + closeable, e);
            }
          });
    }
    this.client = null;
    executorService.shutdown();
  }

  @AutoService(HttpClient.Factory.class)
  @HttpClientName("jdk-http-client")
  public static class Factory implements HttpClient.Factory {

    @Override
    public HttpClient createClient(ClientConfig config) {
      Objects.requireNonNull(config, "Client config must be set");
      return new JdkHttpClient(config);
    }
  }

  private static class HttpProxySelector extends ProxySelector {
    private final Proxy proxy;

    public HttpProxySelector(Proxy proxy) {
      this.proxy = proxy;
    }

    @Override
    public List<Proxy> select(URI uri) {
      if (proxy == null) {
        return List.of();
      }
      if (uri.getScheme().toLowerCase(Locale.ENGLISH).startsWith("http")) {
        return List.of(proxy);
      }
      return List.of();
    }

    @Override
    public void connectFailed(URI uri, SocketAddress sa, IOException ioe) {
      // Do nothing
    }
  }
}
