// 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.router;

import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
import static org.assertj.core.api.Assertions.assertThat;
import static org.openqa.selenium.json.Json.JSON_UTF_8;
import static org.openqa.selenium.remote.http.Contents.asJson;
import static org.openqa.selenium.remote.http.HttpMethod.POST;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.net.URI;
import java.net.URISyntaxException;
import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.SessionNotCreatedException;
import org.openqa.selenium.events.EventBus;
import org.openqa.selenium.events.local.GuavaEventBus;
import org.openqa.selenium.grid.config.MapConfig;
import org.openqa.selenium.grid.data.DefaultSlotMatcher;
import org.openqa.selenium.grid.data.Session;
import org.openqa.selenium.grid.distributor.Distributor;
import org.openqa.selenium.grid.distributor.local.LocalDistributor;
import org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector;
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.server.BaseServerOptions;
import org.openqa.selenium.grid.server.Server;
import org.openqa.selenium.grid.sessionmap.SessionMap;
import org.openqa.selenium.grid.sessionmap.local.LocalSessionMap;
import org.openqa.selenium.grid.sessionqueue.NewSessionQueue;
import org.openqa.selenium.grid.sessionqueue.local.LocalNewSessionQueue;
import org.openqa.selenium.grid.testing.PassthroughHttpClient;
import org.openqa.selenium.grid.testing.TestSessionFactory;
import org.openqa.selenium.grid.web.CombinedHandler;
import org.openqa.selenium.grid.web.EnsureSpecCompliantHeaders;
import org.openqa.selenium.net.PortProber;
import org.openqa.selenium.netty.server.NettyServer;
import org.openqa.selenium.remote.http.Contents;
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.http.Routable;
import org.openqa.selenium.remote.tracing.DefaultTestTracer;
import org.openqa.selenium.remote.tracing.Tracer;
import org.openqa.selenium.testing.drivers.Browser;

class NewSessionCreationTest {

  private static final int newSessionThreadPoolSize = Runtime.getRuntime().availableProcessors();
  private Tracer tracer;
  private EventBus bus;
  private HttpClient.Factory clientFactory;
  private Secret registrationSecret;
  private Server<?> server;

  @BeforeEach
  public void setup() {
    tracer = DefaultTestTracer.createTracer();
    bus = new GuavaEventBus();
    clientFactory = HttpClient.Factory.createDefault();
    registrationSecret = new Secret("hereford hop");
  }

  @AfterEach
  public void stopServer() {
    server.stop();
  }

  @Test
  void ensureJsCannotCreateANewSession() throws URISyntaxException {
    SessionMap sessions = new LocalSessionMap(tracer, bus);
    NewSessionQueue queue =
        new LocalNewSessionQueue(
            tracer,
            new DefaultSlotMatcher(),
            Duration.ofSeconds(2),
            Duration.ofSeconds(60),
            registrationSecret,
            5);

    Distributor distributor =
        new LocalDistributor(
            tracer,
            bus,
            clientFactory,
            sessions,
            queue,
            new DefaultSlotSelector(),
            registrationSecret,
            Duration.ofMinutes(5),
            false,
            Duration.ofSeconds(5),
            newSessionThreadPoolSize);

    Routable router =
        new Router(tracer, clientFactory, sessions, queue, distributor)
            .with(new EnsureSpecCompliantHeaders(ImmutableList.of(), ImmutableSet.of()));

    server =
        new NettyServer(
                new BaseServerOptions(new MapConfig(ImmutableMap.of())),
                router,
                new ProxyWebsocketsIntoGrid(clientFactory, sessions))
            .start();

    URI uri = server.getUrl().toURI();
    Node node =
        LocalNode.builder(tracer, bus, uri, uri, registrationSecret)
            .add(
                Browser.detect().getCapabilities(),
                new TestSessionFactory(
                    (id, caps) ->
                        new Session(
                            id, uri, Browser.detect().getCapabilities(), caps, Instant.now())))
            .build();
    distributor.add(node);

    try (HttpClient client = HttpClient.Factory.createDefault().createClient(server.getUrl())) {
      // Attempt to create a session with an origin header but content type set
      HttpResponse res =
          client.execute(
              new HttpRequest(POST, "/session")
                  .addHeader("Content-Type", JSON_UTF_8)
                  .addHeader("Origin", "localhost")
                  .setContent(
                      Contents.asJson(
                          ImmutableMap.of(
                              "capabilities",
                              ImmutableMap.of(
                                  "alwaysMatch", Browser.detect().getCapabilities())))));

      assertThat(res.getStatus()).isEqualTo(HTTP_INTERNAL_ERROR);

      // And now make sure the session is just fine
      res =
          client.execute(
              new HttpRequest(POST, "/session")
                  .addHeader("Content-Type", JSON_UTF_8)
                  .setContent(
                      Contents.asJson(
                          ImmutableMap.of(
                              "capabilities",
                              ImmutableMap.of(
                                  "alwaysMatch", Browser.detect().getCapabilities())))));

      assertThat(res.isSuccessful()).isTrue();
    }
  }

  @Test
  void shouldNotRetryNewSessionRequestOnUnexpectedError() throws URISyntaxException {
    Capabilities capabilities = new ImmutableCapabilities("browserName", "cheese");
    int nodePort = PortProber.findFreePort();
    URI nodeUri = new URI("http://localhost:" + nodePort);
    CombinedHandler handler = new CombinedHandler();

    SessionMap sessions = new LocalSessionMap(tracer, bus);
    handler.addHandler(sessions);
    NewSessionQueue queue =
        new LocalNewSessionQueue(
            tracer,
            new DefaultSlotMatcher(),
            Duration.ofSeconds(2),
            Duration.ofSeconds(10),
            registrationSecret,
            5);
    handler.addHandler(queue);

    AtomicInteger count = new AtomicInteger();

    // First session creation attempt throws an error.
    // Does not reach second attempt.
    TestSessionFactory sessionFactory =
        new TestSessionFactory(
            (id, caps) -> {
              if (count.get() == 0) {
                count.incrementAndGet();
                throw new SessionNotCreatedException("Expected the exception");
              } else {
                return new Session(id, nodeUri, new ImmutableCapabilities(), caps, Instant.now());
              }
            });

    LocalNode localNode =
        LocalNode.builder(tracer, bus, nodeUri, nodeUri, registrationSecret)
            .add(capabilities, sessionFactory)
            .build();
    handler.addHandler(localNode);

    Distributor distributor =
        new LocalDistributor(
            tracer,
            bus,
            new PassthroughHttpClient.Factory(handler),
            sessions,
            queue,
            new DefaultSlotSelector(),
            registrationSecret,
            Duration.ofMinutes(5),
            false,
            Duration.ofSeconds(5),
            newSessionThreadPoolSize);
    handler.addHandler(distributor);

    distributor.add(localNode);

    Router router = new Router(tracer, clientFactory, sessions, queue, distributor);
    handler.addHandler(router);

    server = new NettyServer(new BaseServerOptions(new MapConfig(ImmutableMap.of())), handler);

    server.start();

    HttpRequest request = new HttpRequest(POST, "/session");
    request.setContent(
        asJson(ImmutableMap.of("capabilities", ImmutableMap.of("alwaysMatch", capabilities))));

    HttpClient client = clientFactory.createClient(server.getUrl());
    HttpResponse httpResponse = client.execute(request);
    assertThat(httpResponse.getStatus()).isEqualTo(HTTP_INTERNAL_ERROR);
  }

  @Test
  @Timeout(10)
  void shouldRejectRequestForUnsupportedCaps() throws URISyntaxException {
    Capabilities capabilities = new ImmutableCapabilities("browserName", "cheese");
    int nodePort = PortProber.findFreePort();
    URI nodeUri = new URI("http://localhost:" + nodePort);
    CombinedHandler handler = new CombinedHandler();

    SessionMap sessions = new LocalSessionMap(tracer, bus);
    handler.addHandler(sessions);
    NewSessionQueue queue =
        new LocalNewSessionQueue(
            tracer,
            new DefaultSlotMatcher(),
            Duration.ofSeconds(5),
            Duration.ofSeconds(60),
            registrationSecret,
            5);
    handler.addHandler(queue);

    TestSessionFactory sessionFactory =
        new TestSessionFactory(
            (id, caps) ->
                new Session(id, nodeUri, new ImmutableCapabilities(), caps, Instant.now()));

    LocalNode localNode =
        LocalNode.builder(tracer, bus, nodeUri, nodeUri, registrationSecret)
            .add(capabilities, sessionFactory)
            .build();
    handler.addHandler(localNode);

    Distributor distributor =
        new LocalDistributor(
            tracer,
            bus,
            new PassthroughHttpClient.Factory(handler),
            sessions,
            queue,
            new DefaultSlotSelector(),
            registrationSecret,
            Duration.ofMinutes(5),
            true,
            Duration.ofSeconds(5),
            newSessionThreadPoolSize);
    handler.addHandler(distributor);

    distributor.add(localNode);

    Router router = new Router(tracer, clientFactory, sessions, queue, distributor);
    handler.addHandler(router);

    server = new NettyServer(new BaseServerOptions(new MapConfig(ImmutableMap.of())), handler);

    server.start();

    HttpRequest request = new HttpRequest(POST, "/session");
    request.setContent(
        asJson(
            ImmutableMap.of(
                "capabilities",
                ImmutableMap.of(
                    "alwaysMatch", new ImmutableCapabilities("browserName", "burger")))));

    HttpClient client = clientFactory.createClient(server.getUrl());
    HttpResponse httpResponse = client.execute(request);
    assertThat(httpResponse.getStatus()).isEqualTo(HTTP_INTERNAL_ERROR);
  }
}
