// 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.sessionqueue.local;

import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static org.openqa.selenium.concurrent.ExecutorServices.shutdownGracefully;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.io.Closeable;
import java.io.IOException;
import java.time.Duration;
import java.time.Instant;
import java.util.Deque;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.SessionNotCreatedException;
import org.openqa.selenium.concurrent.GuardedRunnable;
import org.openqa.selenium.grid.config.Config;
import org.openqa.selenium.grid.data.CreateSessionResponse;
import org.openqa.selenium.grid.data.RequestId;
import org.openqa.selenium.grid.data.SessionRequest;
import org.openqa.selenium.grid.data.SessionRequestCapability;
import org.openqa.selenium.grid.data.SlotMatcher;
import org.openqa.selenium.grid.data.TraceSessionRequest;
import org.openqa.selenium.grid.distributor.config.DistributorOptions;
import org.openqa.selenium.grid.jmx.JMXHelper;
import org.openqa.selenium.grid.jmx.ManagedAttribute;
import org.openqa.selenium.grid.jmx.ManagedService;
import org.openqa.selenium.grid.log.LoggingOptions;
import org.openqa.selenium.grid.security.Secret;
import org.openqa.selenium.grid.security.SecretOptions;
import org.openqa.selenium.grid.sessionqueue.NewSessionQueue;
import org.openqa.selenium.grid.sessionqueue.config.NewSessionQueueOptions;
import org.openqa.selenium.internal.Either;
import org.openqa.selenium.internal.Require;
import org.openqa.selenium.remote.http.Contents;
import org.openqa.selenium.remote.http.HttpResponse;
import org.openqa.selenium.remote.tracing.Span;
import org.openqa.selenium.remote.tracing.TraceContext;
import org.openqa.selenium.remote.tracing.Tracer;

/**
 * An in-memory implementation of the list of new session requests.
 *
 * <p>The lifecycle of a request can be described as:
 *
 * <ol>
 *   <li>User adds an item on to the queue using {@link #addToQueue(SessionRequest)}. This will
 *       block until the request completes in some way.
 *   <li>If the session request is completed, then {@link #complete(RequestId, Either)} must be
 *       called. This will ensure that {@link #addToQueue(SessionRequest)} returns.
 *   <li>If the request cannot be handled right now, call {@link #retryAddToQueue(SessionRequest)}
 *       to return the session request to the front of the queue.
 * </ol>
 *
 * <p>There is a background thread that will reap {@link SessionRequest}s that have timed out. This
 * means that a request can either complete by a listener calling {@link #complete(RequestId,
 * Either)} directly, or by being reaped by the thread.
 */
@ManagedService(
    objectName = "org.seleniumhq.grid:type=SessionQueue,name=LocalSessionQueue",
    description = "New session queue")
public class LocalNewSessionQueue extends NewSessionQueue implements Closeable {

  private static final String NAME = "Local New Session Queue";
  private final SlotMatcher slotMatcher;
  private final Duration requestTimeout;
  private final int batchSize;
  private final Map<RequestId, Data> requests;
  private final Map<RequestId, TraceContext> contexts;
  private final Deque<SessionRequest> queue;
  private final ReadWriteLock lock = new ReentrantReadWriteLock();
  private final ScheduledExecutorService service =
      Executors.newSingleThreadScheduledExecutor(
          r -> {
            Thread thread = new Thread(r);
            thread.setDaemon(true);
            thread.setName(NAME);
            return thread;
          });

  public LocalNewSessionQueue(
      Tracer tracer,
      SlotMatcher slotMatcher,
      Duration requestTimeoutCheck,
      Duration requestTimeout,
      Secret registrationSecret,
      int batchSize) {
    super(tracer, registrationSecret);

    this.slotMatcher = Require.nonNull("Slot matcher", slotMatcher);
    Require.nonNegative("Retry period", requestTimeoutCheck);

    this.requestTimeout = Require.positive("Request timeout", requestTimeout);

    this.requests = new ConcurrentHashMap<>();
    this.queue = new ConcurrentLinkedDeque<>();
    this.contexts = new ConcurrentHashMap<>();

    this.batchSize = Require.positive("Batch size", batchSize);

    service.scheduleAtFixedRate(
        GuardedRunnable.guard(this::timeoutSessions),
        requestTimeoutCheck.toMillis(),
        requestTimeoutCheck.toMillis(),
        MILLISECONDS);

    new JMXHelper().register(this);
  }

  public static NewSessionQueue create(Config config) {
    LoggingOptions loggingOptions = new LoggingOptions(config);
    Tracer tracer = loggingOptions.getTracer();

    NewSessionQueueOptions newSessionQueueOptions = new NewSessionQueueOptions(config);
    SecretOptions secretOptions = new SecretOptions(config);
    SlotMatcher slotMatcher = new DistributorOptions(config).getSlotMatcher();

    return new LocalNewSessionQueue(
        tracer,
        slotMatcher,
        newSessionQueueOptions.getSessionRequestTimeoutPeriod(),
        newSessionQueueOptions.getSessionRequestTimeout(),
        secretOptions.getRegistrationSecret(),
        newSessionQueueOptions.getBatchSize());
  }

  private void timeoutSessions() {
    Instant now = Instant.now();

    Lock readLock = lock.readLock();
    readLock.lock();
    Set<RequestId> ids;
    try {
      ids =
          requests.entrySet().stream()
              .filter(
                  entry ->
                      queue.stream()
                          .anyMatch(
                              sessionRequest ->
                                  sessionRequest.getRequestId().equals(entry.getKey())))
              .filter(entry -> isTimedOut(now, entry.getValue()))
              .map(Map.Entry::getKey)
              .collect(ImmutableSet.toImmutableSet());
    } finally {
      readLock.unlock();
    }
    ids.forEach(this::failDueToTimeout);
  }

  private boolean isTimedOut(Instant now, Data data) {
    return data.endTime.isBefore(now);
  }

  @Override
  public boolean peekEmpty() {
    Lock readLock = lock.readLock();
    readLock.lock();
    try {
      return requests.isEmpty() && queue.isEmpty();
    } finally {
      readLock.unlock();
    }
  }

  @Override
  public HttpResponse addToQueue(SessionRequest request) {
    Require.nonNull("New session request", request);
    Require.nonNull("Request id", request.getRequestId());

    TraceContext context = TraceSessionRequest.extract(tracer, request);
    try (Span ignored = context.createSpan("sessionqueue.add_to_queue")) {
      contexts.put(request.getRequestId(), context);
      Data data = injectIntoQueue(request);

      if (isTimedOut(Instant.now(), data)) {
        failDueToTimeout(request.getRequestId());
      }

      Either<SessionNotCreatedException, CreateSessionResponse> result;
      try {

        boolean sessionCreated = data.latch.await(requestTimeout.toMillis(), MILLISECONDS);
        if (!(sessionCreated || isRequestInQueue(request.getRequestId()))) {
          sessionCreated = data.latch.await(5000, MILLISECONDS);
        }

        if (sessionCreated) {
          result = data.result;
        } else {
          result = Either.left(new SessionNotCreatedException("New session request timed out"));
        }
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        result =
            Either.left(new SessionNotCreatedException("Interrupted when creating the session", e));
      } catch (RuntimeException e) {
        result =
            Either.left(
                new SessionNotCreatedException("An error occurred creating the session", e));
      }

      Lock writeLock = this.lock.writeLock();
      writeLock.lock();
      try {
        requests.remove(request.getRequestId());
        queue.remove(request);
      } finally {
        writeLock.unlock();
      }

      HttpResponse res = new HttpResponse();
      if (result.isRight()) {
        res.setContent(Contents.bytes(result.right().getDownstreamEncodedResponse()));
      } else {
        res.setStatus(HTTP_INTERNAL_ERROR)
            .setContent(
                Contents.asJson(
                    ImmutableMap.of(
                        "value",
                        ImmutableMap.of(
                            "error", "session not created",
                            "message", result.left().getMessage(),
                            "stacktrace", result.left().getStackTrace()))));
      }

      return res;
    }
  }

  @VisibleForTesting
  Data injectIntoQueue(SessionRequest request) {
    Require.nonNull("Session request", request);

    Data data = new Data(request.getEnqueued());

    Lock writeLock = lock.writeLock();
    writeLock.lock();
    try {
      requests.put(request.getRequestId(), data);
      queue.addLast(request);
    } finally {
      writeLock.unlock();
    }

    return data;
  }

  @Override
  public boolean retryAddToQueue(SessionRequest request) {
    Require.nonNull("New session request", request);

    boolean added;
    TraceContext context =
        contexts.getOrDefault(request.getRequestId(), tracer.getCurrentContext());
    try (Span ignored = context.createSpan("sessionqueue.retry")) {
      Lock writeLock = lock.writeLock();
      writeLock.lock();
      try {
        if (!requests.containsKey(request.getRequestId())) {
          return false;
        }

        if (queue.contains(request)) {
          // No need to re-add this
          return true;
        } else {
          added = queue.offerFirst(request);
        }
      } finally {
        writeLock.unlock();
      }

      return added;
    }
  }

  @Override
  public Optional<SessionRequest> remove(RequestId reqId) {
    Require.nonNull("Request ID", reqId);

    Lock writeLock = lock.writeLock();
    writeLock.lock();
    try {
      Iterator<SessionRequest> iterator = queue.iterator();
      while (iterator.hasNext()) {
        SessionRequest req = iterator.next();
        if (reqId.equals(req.getRequestId())) {
          iterator.remove();

          return Optional.of(req);
        }
      }
      return Optional.empty();
    } finally {
      writeLock.unlock();
    }
  }

  private boolean isRequestInQueue(RequestId requestId) {
    Lock readLock = lock.readLock();
    readLock.lock();

    try {
      Optional<SessionRequest> result =
          queue.stream().filter(req -> req.getRequestId().equals(requestId)).findAny();
      return result.isPresent();
    } finally {
      readLock.unlock();
    }
  }

  @Override
  public List<SessionRequest> getNextAvailable(Map<Capabilities, Long> stereotypes) {
    Require.nonNull("Stereotypes", stereotypes);

    Predicate<Capabilities> matchesStereotype =
        caps ->
            stereotypes.entrySet().stream()
                .filter(entry -> entry.getValue() > 0)
                .anyMatch(
                    entry -> {
                      boolean matches = slotMatcher.matches(entry.getKey(), caps);
                      if (matches) {
                        Long value = entry.getValue();
                        entry.setValue(value - 1);
                      }
                      return matches;
                    });

    Lock writeLock = lock.writeLock();
    writeLock.lock();
    try {
      List<SessionRequest> availableRequests =
          queue.stream()
              .filter(req -> req.getDesiredCapabilities().stream().anyMatch(matchesStereotype))
              .limit(batchSize)
              .collect(Collectors.toList());

      availableRequests.forEach(req -> this.remove(req.getRequestId()));

      return availableRequests;
    } finally {
      writeLock.unlock();
    }
  }

  @Override
  public void complete(
      RequestId reqId, Either<SessionNotCreatedException, CreateSessionResponse> result) {
    Require.nonNull("New session request", reqId);
    Require.nonNull("Result", result);
    TraceContext context = contexts.getOrDefault(reqId, tracer.getCurrentContext());
    try (Span ignored = context.createSpan("sessionqueue.completed")) {
      Lock readLock = lock.readLock();
      readLock.lock();
      Data data;
      try {
        data = requests.get(reqId);
      } finally {
        readLock.unlock();
      }

      if (data == null) {
        return;
      }

      Lock writeLock = lock.writeLock();
      writeLock.lock();
      try {
        requests.remove(reqId);
        queue.removeIf(req -> reqId.equals(req.getRequestId()));
        contexts.remove(reqId);
      } finally {
        writeLock.unlock();
      }

      data.setResult(result);
    }
  }

  @Override
  public int clearQueue() {
    Lock writeLock = lock.writeLock();
    writeLock.lock();

    try {
      int size = queue.size();
      queue.clear();
      requests.forEach(
          (reqId, data) ->
              data.setResult(
                  Either.left(new SessionNotCreatedException("Request queue was cleared"))));
      requests.clear();
      return size;
    } finally {
      writeLock.unlock();
    }
  }

  @Override
  public List<SessionRequestCapability> getQueueContents() {
    Lock readLock = lock.readLock();
    readLock.lock();

    try {
      return queue.stream()
          .map(
              req -> new SessionRequestCapability(req.getRequestId(), req.getDesiredCapabilities()))
          .collect(Collectors.toList());
    } finally {
      readLock.unlock();
    }
  }

  @ManagedAttribute(name = "NewSessionQueueSize")
  public int getQueueSize() {
    return queue.size();
  }

  @Override
  public boolean isReady() {
    return true;
  }

  @Override
  public void close() throws IOException {
    shutdownGracefully(NAME, service);
  }

  private void failDueToTimeout(RequestId reqId) {
    complete(reqId, Either.left(new SessionNotCreatedException("Timed out creating session")));
  }

  private class Data {

    public final Instant endTime;
    private final CountDownLatch latch = new CountDownLatch(1);
    public Either<SessionNotCreatedException, CreateSessionResponse> result;
    private boolean complete;

    public Data(Instant enqueued) {
      this.endTime = enqueued.plus(requestTimeout);
      this.result = Either.left(new SessionNotCreatedException("Session not created"));
    }

    public synchronized void setResult(
        Either<SessionNotCreatedException, CreateSessionResponse> result) {
      if (complete) {
        return;
      }
      this.result = result;
      complete = true;
      latch.countDown();
    }
  }
}
