// 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.grid.node.docker;

import static java.util.Optional.ofNullable;
import static org.openqa.selenium.docker.ContainerConfig.image;
import static org.openqa.selenium.remote.Dialect.W3C;
import static org.openqa.selenium.remote.http.Contents.string;
import static org.openqa.selenium.remote.http.HttpMethod.GET;
import static org.openqa.selenium.remote.tracing.Tags.EXCEPTION;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URL;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.time.Duration;
import java.time.Instant;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.TimeZone;
import java.util.function.Predicate;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.Dimension;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.PersistentCapabilities;
import org.openqa.selenium.RetrySessionRequestException;
import org.openqa.selenium.SessionNotCreatedException;
import org.openqa.selenium.TimeoutException;
import org.openqa.selenium.WebDriverException;
import org.openqa.selenium.docker.Container;
import org.openqa.selenium.docker.ContainerConfig;
import org.openqa.selenium.docker.ContainerInfo;
import org.openqa.selenium.docker.Device;
import org.openqa.selenium.docker.Docker;
import org.openqa.selenium.docker.Image;
import org.openqa.selenium.docker.Port;
import org.openqa.selenium.grid.data.CreateSessionRequest;
import org.openqa.selenium.grid.node.ActiveSession;
import org.openqa.selenium.grid.node.SessionFactory;
import org.openqa.selenium.internal.Either;
import org.openqa.selenium.internal.Require;
import org.openqa.selenium.json.Json;
import org.openqa.selenium.net.PortProber;
import org.openqa.selenium.remote.Command;
import org.openqa.selenium.remote.Dialect;
import org.openqa.selenium.remote.DriverCommand;
import org.openqa.selenium.remote.ProtocolHandshake;
import org.openqa.selenium.remote.Response;
import org.openqa.selenium.remote.SessionId;
import org.openqa.selenium.remote.http.ClientConfig;
import org.openqa.selenium.remote.http.HttpClient;
import org.openqa.selenium.remote.http.HttpRequest;
import org.openqa.selenium.remote.http.HttpResponse;
import org.openqa.selenium.remote.tracing.AttributeKey;
import org.openqa.selenium.remote.tracing.AttributeMap;
import org.openqa.selenium.remote.tracing.Span;
import org.openqa.selenium.remote.tracing.Status;
import org.openqa.selenium.remote.tracing.Tracer;
import org.openqa.selenium.support.ui.FluentWait;
import org.openqa.selenium.support.ui.Wait;

public class DockerSessionFactory implements SessionFactory {

  private static final Logger LOG = Logger.getLogger(DockerSessionFactory.class.getName());

  private final Tracer tracer;
  private final HttpClient.Factory clientFactory;
  private final Duration sessionTimeout;
  private final Docker docker;
  private final URI dockerUri;
  private final Image browserImage;
  private final Capabilities stereotype;
  private final List<Device> devices;
  private final Image videoImage;
  private final DockerAssetsPath assetsPath;
  private final String networkName;
  private final boolean runningInDocker;
  private final Predicate<Capabilities> predicate;
  private final Map<String, Object> hostConfig;
  private final List<String> hostConfigKeys;

  public DockerSessionFactory(
      Tracer tracer,
      HttpClient.Factory clientFactory,
      Duration sessionTimeout,
      Docker docker,
      URI dockerUri,
      Image browserImage,
      Capabilities stereotype,
      List<Device> devices,
      Image videoImage,
      DockerAssetsPath assetsPath,
      String networkName,
      boolean runningInDocker,
      Predicate<Capabilities> predicate,
      Map<String, Object> hostConfig,
      List<String> hostConfigKeys) {
    this.tracer = Require.nonNull("Tracer", tracer);
    this.clientFactory = Require.nonNull("HTTP client", clientFactory);
    this.sessionTimeout = Require.nonNull("Session timeout", sessionTimeout);
    this.docker = Require.nonNull("Docker command", docker);
    this.dockerUri = Require.nonNull("Docker URI", dockerUri);
    this.browserImage = Require.nonNull("Docker browser image", browserImage);
    this.networkName = Require.nonNull("Docker network name", networkName);
    this.stereotype = ImmutableCapabilities.copyOf(Require.nonNull("Stereotype", stereotype));
    this.devices = Require.nonNull("Container devices", devices);
    this.videoImage = videoImage;
    this.assetsPath = assetsPath;
    this.runningInDocker = runningInDocker;
    this.predicate = Require.nonNull("Accepted capabilities predicate", predicate);
    this.hostConfig = Require.nonNull("Container host config", hostConfig);
    this.hostConfigKeys = Require.nonNull("Browser container host config keys", hostConfigKeys);
  }

  @Override
  public Capabilities getStereotype() {
    return stereotype;
  }

  @Override
  public boolean test(Capabilities capabilities) {
    return predicate.test(capabilities);
  }

  @Override
  public Either<WebDriverException, ActiveSession> apply(CreateSessionRequest sessionRequest) {
    LOG.info("Starting session for " + sessionRequest.getDesiredCapabilities());

    int port = runningInDocker ? 4444 : PortProber.findFreePort();
    try (Span span = tracer.getCurrentContext().createSpan("docker_session_factory.apply")) {
      AttributeMap attributeMap = tracer.createAttributeMap();
      attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), this.getClass().getName());
      String logMessage =
          runningInDocker
              ? "Creating container..."
              : "Creating container, mapping container port 4444 to " + port;
      LOG.info(logMessage);
      Container container = createBrowserContainer(port, sessionRequest.getDesiredCapabilities());
      container.start();
      ContainerInfo containerInfo = container.inspect();

      String containerIp = containerInfo.getIp();
      URL remoteAddress = getUrl(port, containerIp);
      ClientConfig clientConfig =
          ClientConfig.defaultConfig().baseUrl(remoteAddress).readTimeout(sessionTimeout);
      HttpClient client = clientFactory.createClient(clientConfig);

      attributeMap.put("docker.browser.image", browserImage.toString());
      attributeMap.put("container.port", port);
      attributeMap.put("container.id", container.getId().toString());
      attributeMap.put("container.ip", containerIp);
      attributeMap.put("docker.server.url", remoteAddress.toString());

      LOG.info(
          String.format(
              "Waiting for server to start (container id: %s, url %s)",
              container.getId(), remoteAddress));
      try {
        waitForServerToStart(client, Duration.ofMinutes(1));
      } catch (TimeoutException e) {
        span.setAttribute(AttributeKey.ERROR.getKey(), true);
        span.setStatus(Status.CANCELLED);

        EXCEPTION.accept(attributeMap, e);
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(),
            "Unable to connect to docker server. Stopping container: " + e.getMessage());
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);

        container.stop(Duration.ofMinutes(1));
        String message =
            String.format(
                "Unable to connect to docker server (container id: %s)", container.getId());
        LOG.warning(message);
        client.close();
        return Either.left(new RetrySessionRequestException(message));
      }
      LOG.info(String.format("Server is ready (container id: %s)", container.getId()));

      Command command =
          new Command(null, DriverCommand.NEW_SESSION(sessionRequest.getDesiredCapabilities()));
      ProtocolHandshake.Result result;
      Response response;
      try {
        result = new ProtocolHandshake().createSession(client, command);
        response = result.createResponse();
        attributeMap.put(AttributeKey.DRIVER_RESPONSE.getKey(), response.toString());
      } catch (IOException | RuntimeException e) {
        span.setAttribute(AttributeKey.ERROR.getKey(), true);
        span.setStatus(Status.CANCELLED);

        EXCEPTION.accept(attributeMap, e);
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(),
            "Unable to create session. Stopping and  container: " + e.getMessage());
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);

        container.stop(Duration.ofMinutes(1));
        String message = "Unable to create session: " + e.getMessage();
        LOG.log(Level.WARNING, message, e);
        client.close();
        return Either.left(new SessionNotCreatedException(message));
      }

      SessionId id = new SessionId(response.getSessionId());
      Capabilities capabilities = new ImmutableCapabilities((Map<?, ?>) response.getValue());
      Capabilities mergedCapabilities = sessionRequest.getDesiredCapabilities().merge(capabilities);
      mergedCapabilities =
          addForwardCdpEndpoint(mergedCapabilities, containerIp, port, id.toString());

      Container videoContainer = null;
      Optional<DockerAssetsPath> path = ofNullable(this.assetsPath);
      if (path.isPresent()) {
        // Seems we can store session assets
        String containerPath = path.get().getContainerPath(id);
        saveSessionCapabilities(mergedCapabilities, containerPath);
        String hostPath = path.get().getHostPath(id);
        videoContainer = startVideoContainer(mergedCapabilities, containerIp, hostPath);
      }

      Dialect downstream =
          sessionRequest.getDownstreamDialects().contains(result.getDialect())
              ? result.getDialect()
              : W3C;
      attributeMap.put(AttributeKey.DOWNSTREAM_DIALECT.getKey(), downstream.toString());
      attributeMap.put(AttributeKey.DRIVER_RESPONSE.getKey(), response.toString());

      span.addEvent("Docker driver service created session", attributeMap);
      LOG.fine(
          String.format(
              "Created session: %s - %s (container id: %s)",
              id, mergedCapabilities, container.getId()));
      return Either.right(
          new DockerSession(
              container,
              videoContainer,
              tracer,
              client,
              id,
              remoteAddress,
              stereotype,
              mergedCapabilities,
              downstream,
              result.getDialect(),
              Instant.now(),
              assetsPath));
    }
  }

  private Capabilities addForwardCdpEndpoint(
      Capabilities sessionCapabilities, String containerIp, int port, String sessionId) {
    // We add this endpoint to go around the situation where a user wants to do CDP over
    // Dynamic Grid. In a conventional Grid setup, this is not needed because the browser will
    // be running on the same host where the Node is running. However, in Dynamic Grid, the Docker
    // Node is running on a different host/container. Therefore, we need to forward the websocket
    // connection to the container where the actual browser is running.
    String forwardCdpPath =
        String.format("ws://%s:%s/session/%s/se/fwd", containerIp, port, sessionId);
    return new PersistentCapabilities(sessionCapabilities)
        .setCapability("se:forwardCdp", forwardCdpPath);
  }

  private Container createBrowserContainer(int port, Capabilities sessionCapabilities) {
    Map<String, String> browserContainerEnvVars = getBrowserContainerEnvVars(sessionCapabilities);
    long browserContainerShmMemorySize = 2147483648L; // 2GB
    ContainerConfig containerConfig =
        image(browserImage)
            .env(browserContainerEnvVars)
            .shmMemorySize(browserContainerShmMemorySize)
            .network(networkName)
            .devices(devices)
            .applyHostConfig(hostConfig, hostConfigKeys);
    if (!runningInDocker) {
      containerConfig = containerConfig.map(Port.tcp(4444), Port.tcp(port));
    }
    LOG.fine("Container config: " + containerConfig);
    return docker.create(containerConfig);
  }

  private Map<String, String> getBrowserContainerEnvVars(Capabilities sessionRequestCapabilities) {
    Map<String, String> envVars = new HashMap<>();
    // Passing env vars set to the child container
    setEnvVarsToContainer(envVars);
    // Capabilities set to env vars with higher precedence
    setCapsToEnvVars(sessionRequestCapabilities, envVars);
    return envVars;
  }

  private void setEnvVarsToContainer(Map<String, String> envVars) {
    Map<String, String> seEnvVars = System.getenv();
    seEnvVars.entrySet().stream()
        .filter(
            entry ->
                entry.getKey().startsWith("SE_") || entry.getKey().equalsIgnoreCase("LANGUAGE"))
        .forEach(entry -> envVars.put(entry.getKey(), entry.getValue()));
  }

  private void setCapsToEnvVars(
      Capabilities sessionRequestCapabilities, Map<String, String> envVars) {
    Optional<Dimension> screenResolution =
        ofNullable(getScreenResolution(sessionRequestCapabilities));
    screenResolution.ifPresent(
        dimension -> {
          envVars.put("SE_SCREEN_WIDTH", String.valueOf(dimension.getWidth()));
          envVars.put("SE_SCREEN_HEIGHT", String.valueOf(dimension.getHeight()));
        });
    Optional<TimeZone> timeZone = ofNullable(getTimeZone(sessionRequestCapabilities));
    timeZone.ifPresent(zone -> envVars.put("TZ", zone.getID()));
  }

  private Container startVideoContainer(
      Capabilities sessionCapabilities, String browserContainerIp, String hostPath) {
    if (!recordVideoForSession(sessionCapabilities)) {
      return null;
    }
    int videoPort = 9000;
    Map<String, String> envVars = getVideoContainerEnvVars(sessionCapabilities, browserContainerIp);
    Map<String, String> volumeBinds = Collections.singletonMap(hostPath, "/videos");
    ContainerConfig containerConfig =
        image(videoImage).env(envVars).bind(volumeBinds).network(networkName);
    if (!runningInDocker) {
      videoPort = PortProber.findFreePort();
      containerConfig = containerConfig.map(Port.tcp(9000), Port.tcp(videoPort));
    }
    Container videoContainer = docker.create(containerConfig);
    videoContainer.start();
    String videoContainerIp = runningInDocker ? videoContainer.inspect().getIp() : "localhost";
    URI videoContainerUrl = URI.create(String.format("http://%s:%s", videoContainerIp, videoPort));
    HttpClient videoClient =
        clientFactory.createClient(ClientConfig.defaultConfig().baseUri(videoContainerUrl));
    try {
      LOG.fine(String.format("Waiting for video recording... (id: %s)", videoContainer.getId()));
      waitForServerToStart(videoClient, Duration.ofMinutes(1));
    } catch (Exception e) {
      videoContainer.stop(Duration.ofSeconds(10));
      String message =
          String.format(
              "Unable to verify video recording started (container id: %s), %s",
              videoContainer.getId(), e.getMessage());
      LOG.warning(message);
      videoClient.close();
      return null;
    }
    LOG.info(String.format("Video container started (id: %s)", videoContainer.getId()));
    return videoContainer;
  }

  private Map<String, String> getVideoContainerEnvVars(
      Capabilities sessionRequestCapabilities, String containerIp) {
    Map<String, String> envVars = new HashMap<>();
    // Passing env vars set to the child container
    setEnvVarsToContainer(envVars);
    // Capabilities set to env vars with higher precedence
    setCapsToEnvVars(sessionRequestCapabilities, envVars);
    envVars.put("DISPLAY_CONTAINER_NAME", containerIp);
    Optional<String> videoName =
        ofNullable(getVideoFileName(sessionRequestCapabilities, "se:videoName"))
            .or(() -> ofNullable(getVideoFileName(sessionRequestCapabilities, "se:name")));
    videoName.ifPresent(name -> envVars.put("SE_VIDEO_FILE_NAME", String.format("%s.mp4", name)));
    return envVars;
  }

  private String getVideoFileName(Capabilities sessionRequestCapabilities, String capabilityName) {
    Optional<Object> testName =
        ofNullable(sessionRequestCapabilities.getCapability(capabilityName));
    if (testName.isPresent()) {
      String name = testName.get().toString();
      if (!name.isEmpty()) {
        name = name.replaceAll(" ", "_").replaceAll("[^a-zA-Z0-9_-]", "");
        if (name.length() > 251) {
          name = name.substring(0, 251);
        }
        return name;
      }
    }
    return null;
  }

  private TimeZone getTimeZone(Capabilities sessionRequestCapabilities) {
    Optional<Object> timeZone = ofNullable(sessionRequestCapabilities.getCapability("se:timeZone"));
    if (timeZone.isPresent()) {
      String tz = timeZone.get().toString();
      if (Arrays.asList(TimeZone.getAvailableIDs()).contains(tz)) {
        return TimeZone.getTimeZone(tz);
      }
    }
    String envTz = System.getenv("TZ");
    if (Arrays.asList(TimeZone.getAvailableIDs()).contains(envTz)) {
      return TimeZone.getTimeZone(envTz);
    }
    return null;
  }

  private Dimension getScreenResolution(Capabilities sessionRequestCapabilities) {
    Optional<Object> screenResolution =
        ofNullable(sessionRequestCapabilities.getCapability("se:screenResolution"));
    if (!screenResolution.isPresent()) {
      return null;
    }
    try {
      String[] resolution = screenResolution.get().toString().split("x");
      int screenWidth = Integer.parseInt(resolution[0]);
      int screenHeight = Integer.parseInt(resolution[1]);
      if (screenWidth > 0 && screenHeight > 0) {
        return new Dimension(screenWidth, screenHeight);
      } else {
        LOG.warning(
            "One of the values provided for screenResolution is negative, "
                + "defaults will be used. Received value: "
                + screenResolution);
      }
    } catch (Exception e) {
      LOG.warning(
          "Values provided for screenResolution are not valid integers or "
              + "either width or height are missing, defaults will be used."
              + "Received value: "
              + screenResolution);
    }
    return null;
  }

  private boolean recordVideoForSession(Capabilities sessionRequestCapabilities) {
    Optional<Object> recordVideo =
        ofNullable(sessionRequestCapabilities.getCapability("se:recordVideo"));
    return recordVideo.isPresent() && Boolean.parseBoolean(recordVideo.get().toString());
  }

  private void saveSessionCapabilities(Capabilities sessionRequestCapabilities, String path) {
    String capsToJson = new Json().toJson(sessionRequestCapabilities);
    try {
      Files.createDirectories(Paths.get(path));
      Files.write(
          Paths.get(path, "sessionCapabilities.json"),
          capsToJson.getBytes(Charset.defaultCharset()));
    } catch (IOException e) {
      LOG.log(Level.WARNING, "Failed to save session capabilities", e);
    }
  }

  private void waitForServerToStart(HttpClient client, Duration duration) {
    Wait<Object> wait =
        new FluentWait<>(new Object()).withTimeout(duration).ignoring(UncheckedIOException.class);

    wait.until(
        obj -> {
          HttpResponse response = client.execute(new HttpRequest(GET, "/status"));
          LOG.fine(string(response));
          return 200 == response.getStatus();
        });
  }

  private URL getUrl(int port, String containerIp) {
    try {
      String host = "localhost";
      if (runningInDocker) {
        host = containerIp;
      } else {
        if (dockerUri.getScheme().startsWith("tcp") || dockerUri.getScheme().startsWith("http")) {
          host = dockerUri.getHost();
        }
      }
      return new URL(String.format("http://%s:%s/wd/hub", host, port));
    } catch (MalformedURLException e) {
      throw new SessionNotCreatedException(e.getMessage(), e);
    }
  }
}
