// 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.distributor.selector;

import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static org.assertj.core.api.Assertions.assertThat;
import static org.openqa.selenium.grid.data.Availability.UP;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.io.UncheckedIOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.time.Duration;
import java.time.Instant;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.UUID;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.events.EventBus;
import org.openqa.selenium.events.local.GuavaEventBus;
import org.openqa.selenium.grid.data.DefaultSlotMatcher;
import org.openqa.selenium.grid.data.NodeId;
import org.openqa.selenium.grid.data.NodeStatus;
import org.openqa.selenium.grid.data.Session;
import org.openqa.selenium.grid.data.Slot;
import org.openqa.selenium.grid.data.SlotId;
import org.openqa.selenium.grid.node.Node;
import org.openqa.selenium.grid.node.local.LocalNode;
import org.openqa.selenium.grid.security.Secret;
import org.openqa.selenium.grid.testing.TestSessionFactory;
import org.openqa.selenium.remote.SessionId;
import org.openqa.selenium.remote.http.HttpHandler;
import org.openqa.selenium.remote.http.HttpRequest;
import org.openqa.selenium.remote.http.HttpResponse;
import org.openqa.selenium.remote.tracing.DefaultTestTracer;
import org.openqa.selenium.remote.tracing.Tracer;

class DefaultSlotSelectorTest {

  private final Random random = new Random();
  private final DefaultSlotSelector selector = new DefaultSlotSelector();
  private Tracer tracer;
  private EventBus bus;
  private URI uri;

  @BeforeEach
  public void setUp() throws URISyntaxException {
    tracer = DefaultTestTracer.createTracer();
    bus = new GuavaEventBus();
    uri = new URI("http://localhost:1234");
  }

  @Test
  void numberOfSupportedBrowsersByNodeIsCorrect() {
    NodeStatus nodeStatus = createNode("edge", "firefox", "chrome");
    long supportedBrowsersByNode = selector.getNumberOfSupportedBrowsers(nodeStatus);
    assertThat(supportedBrowsersByNode).isEqualTo(3);
    nodeStatus = createNode("edge", "chrome");
    supportedBrowsersByNode = selector.getNumberOfSupportedBrowsers(nodeStatus);
    assertThat(supportedBrowsersByNode).isEqualTo(2);
    nodeStatus = createNode("chrome");
    supportedBrowsersByNode = selector.getNumberOfSupportedBrowsers(nodeStatus);
    assertThat(supportedBrowsersByNode).isEqualTo(1);
  }

  @Test
  void nodesAreOrderedNodesByBrowserVersion() {
    Capabilities caps = new ImmutableCapabilities("browserName", "chrome");

    NodeStatus node1 =
        createNodeWithStereotypes(
            Arrays.asList(
                ImmutableMap.of("browserName", "chrome", "browserVersion", "131.0"),
                ImmutableMap.of("browserName", "chrome", "browserVersion", "132.0")));
    NodeStatus node2 =
        createNodeWithStereotypes(
            Arrays.asList(ImmutableMap.of("browserName", "chrome", "browserVersion", "131.0")));
    NodeStatus node3 =
        createNodeWithStereotypes(
            Arrays.asList(ImmutableMap.of("browserName", "chrome", "browserVersion", "")));
    NodeStatus node4 =
        createNodeWithStereotypes(
            Arrays.asList(ImmutableMap.of("browserName", "chrome", "browserVersion", "131.1")));
    NodeStatus node5 =
        createNodeWithStereotypes(
            Arrays.asList(ImmutableMap.of("browserName", "chrome", "browserVersion", "beta")));
    Set<NodeStatus> nodes = ImmutableSet.of(node1, node2, node3, node4, node5);

    Set<SlotId> slots = selector.selectSlot(caps, nodes, new DefaultSlotMatcher());

    ImmutableSet<NodeId> nodeIds =
        slots.stream().map(SlotId::getOwningNodeId).distinct().collect(toImmutableSet());

    assertThat(nodeIds)
        .containsSequence(
            node3.getNodeId(),
            node1.getNodeId(),
            node4.getNodeId(),
            node2.getNodeId(),
            node5.getNodeId());
  }

  @Test
  void nodesAreOrderedNodesByNumberOfSupportedBrowsers() {
    Set<NodeStatus> nodes = new HashSet<>();

    Capabilities caps = new ImmutableCapabilities("browserName", "chrome");

    NodeStatus threeBrowsers = createNode("edge", "firefox", "chrome");
    NodeStatus twoBrowsers = createNode("firefox", "chrome");
    NodeStatus oneBrowser = createNode("chrome");
    nodes.add(threeBrowsers);
    nodes.add(twoBrowsers);
    nodes.add(oneBrowser);

    Set<SlotId> slots = selector.selectSlot(caps, nodes, new DefaultSlotMatcher());

    ImmutableSet<NodeId> nodeIds =
        slots.stream().map(SlotId::getOwningNodeId).distinct().collect(toImmutableSet());

    assertThat(nodeIds)
        .containsSequence(
            oneBrowser.getNodeId(), twoBrowsers.getNodeId(), threeBrowsers.getNodeId());
  }

  @Test
  void theMostLightlyLoadedNodeIsSelectedFirst() {
    // Create enough hosts so that we avoid the scheduler returning hosts in:
    // * insertion order
    // * reverse insertion order
    // * sorted with most heavily used first

    Capabilities caps = new ImmutableCapabilities("cheese", "beyaz peynir");

    NodeStatus lightest = createNode(Collections.singletonList(caps), 10, 0);
    NodeStatus medium = createNode(Collections.singletonList(caps), 10, 4);
    NodeStatus heavy = createNode(Collections.singletonList(caps), 10, 6);
    NodeStatus massive = createNode(Collections.singletonList(caps), 10, 8);

    Set<SlotId> ids =
        selector.selectSlot(
            caps, ImmutableSet.of(heavy, medium, lightest, massive), new DefaultSlotMatcher());
    SlotId expected = ids.iterator().next();

    assertThat(lightest.getSlots().stream()).anyMatch(slot -> expected.equals(slot.getId()));
  }

  @Test
  void theNodeWhichHasExceededMaxSessionsIsNotSelected() {
    Capabilities chrome = new ImmutableCapabilities("browserName", "chrome");

    NodeStatus lightLoad = createNode(ImmutableList.of(chrome), 12, 2);
    NodeStatus mediumLoad = createNode(ImmutableList.of(chrome), 12, 5);
    NodeStatus maximumLoad = createNode(ImmutableList.of(chrome), 12, 12);

    Set<SlotId> ids =
        selector.selectSlot(
            chrome, ImmutableSet.of(maximumLoad, mediumLoad, lightLoad), new DefaultSlotMatcher());
    SlotId expected = ids.iterator().next();

    // The slot should belong to the Node with light load
    assertThat(lightLoad.getSlots().stream()).anyMatch(slot -> expected.equals(slot.getId()));

    // The node whose current number of sessions is greater than or equal to the max sessions is not
    // included
    // Hence, the node with the maximum load is skipped
    ImmutableSet<NodeId> nodeIds =
        ids.stream().map(SlotId::getOwningNodeId).distinct().collect(toImmutableSet());
    assertThat(nodeIds).doesNotContain(maximumLoad.getNodeId());
    assertThat(nodeIds).containsSequence(lightLoad.getNodeId(), mediumLoad.getNodeId());
  }

  @Test
  void nodesAreOrderedByNumberOfSupportedBrowsersAndLoad() {
    Capabilities chrome = new ImmutableCapabilities("browserName", "chrome");
    Capabilities firefox = new ImmutableCapabilities("browserName", "firefox");
    Capabilities safari = new ImmutableCapabilities("browserName", "safari");

    NodeStatus lightLoadAndThreeBrowsers =
        createNode(ImmutableList.of(chrome, firefox, safari), 12, 2);
    NodeStatus mediumLoadAndTwoBrowsers = createNode(ImmutableList.of(chrome, firefox), 12, 5);
    NodeStatus mediumLoadAndOtherTwoBrowsers = createNode(ImmutableList.of(safari, chrome), 12, 6);
    NodeStatus highLoadAndOneBrowser = createNode(ImmutableList.of(chrome), 12, 8);

    Set<SlotId> ids =
        selector.selectSlot(
            chrome,
            ImmutableSet.of(
                lightLoadAndThreeBrowsers,
                mediumLoadAndTwoBrowsers,
                mediumLoadAndOtherTwoBrowsers,
                highLoadAndOneBrowser),
            new DefaultSlotMatcher());

    // The slot should belong to the Node with high load because it only supports Chrome, leaving
    // the other Nodes with more availability for other browsers
    SlotId expected = ids.iterator().next();
    assertThat(highLoadAndOneBrowser.getSlots().stream())
        .anyMatch(slot -> expected.equals(slot.getId()));

    // Nodes are ordered by the diversity of supported browsers, then by load
    // The node whose current number of sessions is greater than or equal to the max sessions is not
    // included
    ImmutableSet<NodeId> nodeIds =
        ids.stream().map(SlotId::getOwningNodeId).distinct().collect(toImmutableSet());
    assertThat(nodeIds)
        .containsSequence(
            highLoadAndOneBrowser.getNodeId(),
            mediumLoadAndTwoBrowsers.getNodeId(),
            lightLoadAndThreeBrowsers.getNodeId());
  }

  private NodeStatus createNode(List<Capabilities> stereotypes, int count, int currentLoad) {
    NodeId nodeId = new NodeId(UUID.randomUUID());

    URI uri = createUri();

    Set<Slot> slots = new HashSet<>();
    stereotypes.forEach(
        stereotype -> {
          for (int i = 0; i < currentLoad; i++) {
            Instant now = Instant.now();
            slots.add(
                new Slot(
                    new SlotId(nodeId, UUID.randomUUID()),
                    stereotype,
                    now,
                    new Session(
                        new SessionId(UUID.randomUUID()), uri, stereotype, stereotype, now)));
          }

          for (int i = 0; i < count - currentLoad; i++) {
            slots.add(
                new Slot(new SlotId(nodeId, UUID.randomUUID()), stereotype, Instant.EPOCH, null));
          }
        });

    return new NodeStatus(
        nodeId,
        uri,
        count,
        ImmutableSet.copyOf(slots),
        UP,
        Duration.ofSeconds(10),
        Duration.ofSeconds(300),
        "4.0.0",
        ImmutableMap.of(
            "name", "Max OS X",
            "arch", "x86_64",
            "version", "10.15.7"));
  }

  private NodeStatus createNode(String... browsers) {
    URI uri = createUri();
    LocalNode.Builder nodeBuilder =
        LocalNode.builder(tracer, bus, uri, uri, new Secret("cornish yarg"));
    nodeBuilder.maximumConcurrentSessions(browsers.length);

    Arrays.stream(browsers)
        .forEach(
            browser -> {
              Capabilities caps = new ImmutableCapabilities("browserName", browser);
              nodeBuilder.add(caps, new TestSessionFactory((id, c) -> new Handler(c)));
            });

    Node myNode = nodeBuilder.build();
    return myNode.getStatus();
  }

  private NodeStatus createNodeWithStereotypes(List<ImmutableMap> stereotypes) {
    URI uri = createUri();
    LocalNode.Builder nodeBuilder =
        LocalNode.builder(tracer, bus, uri, uri, new Secret("cornish yarg"));
    nodeBuilder.maximumConcurrentSessions(stereotypes.size());
    stereotypes.forEach(
        stereotype -> {
          Capabilities caps = new ImmutableCapabilities(stereotype);
          nodeBuilder.add(caps, new TestSessionFactory((id, c) -> new Handler(c)));
        });
    Node myNode = nodeBuilder.build();
    return myNode.getStatus();
  }

  private URI createUri() {
    try {
      return new URI("http://localhost:" + random.nextInt());
    } catch (URISyntaxException e) {
      throw new RuntimeException(e);
    }
  }

  private class Handler extends Session implements HttpHandler {
    private Handler(Capabilities capabilities) {
      super(
          new SessionId(UUID.randomUUID()),
          uri,
          new ImmutableCapabilities(),
          capabilities,
          Instant.now());
    }

    @Override
    public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
      return new HttpResponse();
    }
  }
}
