/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF 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 com.datastax.oss.driver.internal.core.cql;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverConfig;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.cql.Statement;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metrics.SessionMetric;
import com.datastax.oss.driver.api.core.retry.RetryPolicy;
import com.datastax.oss.driver.api.core.session.Request;
import com.datastax.oss.driver.api.core.session.Session;
import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy;
import com.datastax.oss.driver.api.core.time.TimestampGenerator;
import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator;
import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
import com.datastax.oss.driver.internal.core.DefaultConsistencyLevelRegistry;
import com.datastax.oss.driver.internal.core.ProtocolFeature;
import com.datastax.oss.driver.internal.core.ProtocolVersionRegistry;
import com.datastax.oss.driver.internal.core.channel.DriverChannel;
import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
import com.datastax.oss.driver.internal.core.context.NettyOptions;
import com.datastax.oss.driver.internal.core.metadata.DefaultMetadata;
import com.datastax.oss.driver.internal.core.metadata.LoadBalancingPolicyWrapper;
import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater;
import com.datastax.oss.driver.internal.core.pool.ChannelPool;
import com.datastax.oss.driver.internal.core.servererrors.DefaultWriteTypeRegistry;
import com.datastax.oss.driver.internal.core.session.DefaultSession;
import com.datastax.oss.driver.internal.core.session.throttling.PassThroughRequestThrottler;
import com.datastax.oss.driver.internal.core.tracker.NoopRequestTracker;
import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer;
import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer.CapturedTimeout;
import com.datastax.oss.protocol.internal.Frame;
import io.netty.channel.EventLoopGroup;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.mockito.stubbing.OngoingStubbing;

/**
 * Provides the environment to test a request handler, where a query plan can be defined, and the
 * behavior of each successive node simulated.
 */
public class RequestHandlerTestHarness implements AutoCloseable {

  public static Builder builder() {
    return new Builder();
  }

  private final CapturingTimer timer = new CapturingTimer();
  private final Map<Node, ChannelPool> pools;

  @Mock protected InternalDriverContext context;
  @Mock protected DefaultSession session;
  @Mock protected EventLoopGroup eventLoopGroup;
  @Mock protected NettyOptions nettyOptions;
  @Mock protected DriverConfig config;
  @Mock protected DriverExecutionProfile defaultProfile;
  @Mock protected LoadBalancingPolicyWrapper loadBalancingPolicyWrapper;
  @Mock protected RetryPolicy retryPolicy;
  @Mock protected SpeculativeExecutionPolicy speculativeExecutionPolicy;
  @Mock protected TimestampGenerator timestampGenerator;
  @Mock protected ProtocolVersionRegistry protocolVersionRegistry;
  @Mock protected SessionMetricUpdater sessionMetricUpdater;

  protected RequestHandlerTestHarness(Builder builder) {
    MockitoAnnotations.initMocks(this);

    when(nettyOptions.getTimer()).thenReturn(timer);
    when(nettyOptions.ioEventLoopGroup()).thenReturn(eventLoopGroup);
    when(context.getNettyOptions()).thenReturn(nettyOptions);

    when(defaultProfile.getName()).thenReturn(DriverExecutionProfile.DEFAULT_NAME);
    // TODO make configurable in the test, also handle profiles
    when(defaultProfile.getDuration(DefaultDriverOption.REQUEST_TIMEOUT))
        .thenReturn(Duration.ofMillis(500));
    when(defaultProfile.getString(DefaultDriverOption.REQUEST_CONSISTENCY))
        .thenReturn(DefaultConsistencyLevel.LOCAL_ONE.name());
    when(defaultProfile.getInt(DefaultDriverOption.REQUEST_PAGE_SIZE)).thenReturn(5000);
    when(defaultProfile.getString(DefaultDriverOption.REQUEST_SERIAL_CONSISTENCY))
        .thenReturn(DefaultConsistencyLevel.SERIAL.name());
    when(defaultProfile.getBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE))
        .thenReturn(builder.defaultIdempotence);
    when(defaultProfile.getBoolean(DefaultDriverOption.PREPARE_ON_ALL_NODES)).thenReturn(true);

    when(config.getDefaultProfile()).thenReturn(defaultProfile);
    when(context.getConfig()).thenReturn(config);

    when(loadBalancingPolicyWrapper.newQueryPlan(
            any(Request.class), anyString(), any(Session.class)))
        .thenReturn(builder.buildQueryPlan());
    when(context.getLoadBalancingPolicyWrapper()).thenReturn(loadBalancingPolicyWrapper);

    when(context.getRetryPolicy(anyString())).thenReturn(retryPolicy);

    // Disable speculative executions by default
    when(speculativeExecutionPolicy.nextExecution(
            any(Node.class), any(CqlIdentifier.class), any(Request.class), anyInt()))
        .thenReturn(-1L);
    when(context.getSpeculativeExecutionPolicy(anyString())).thenReturn(speculativeExecutionPolicy);

    when(context.getCodecRegistry()).thenReturn(CodecRegistry.DEFAULT);

    when(timestampGenerator.next()).thenReturn(Statement.NO_DEFAULT_TIMESTAMP);
    when(context.getTimestampGenerator()).thenReturn(timestampGenerator);

    pools = builder.buildMockPools();
    when(session.getChannel(any(Node.class), anyString()))
        .thenAnswer(
            invocation -> {
              Node node = invocation.getArgument(0);
              return pools.get(node).next();
            });
    when(session.getRepreparePayloads()).thenReturn(new ConcurrentHashMap<>());

    when(session.setKeyspace(any(CqlIdentifier.class)))
        .thenReturn(CompletableFuture.completedFuture(null));

    when(session.getMetricUpdater()).thenReturn(sessionMetricUpdater);
    when(sessionMetricUpdater.isEnabled(any(SessionMetric.class), anyString())).thenReturn(true);

    when(session.getMetadata()).thenReturn(DefaultMetadata.EMPTY);

    when(context.getProtocolVersionRegistry()).thenReturn(protocolVersionRegistry);
    when(protocolVersionRegistry.supports(any(ProtocolVersion.class), any(ProtocolFeature.class)))
        .thenReturn(true);

    if (builder.protocolVersion != null) {
      when(context.getProtocolVersion()).thenReturn(builder.protocolVersion);
    }

    when(context.getConsistencyLevelRegistry()).thenReturn(new DefaultConsistencyLevelRegistry());

    when(context.getWriteTypeRegistry()).thenReturn(new DefaultWriteTypeRegistry());

    when(context.getRequestThrottler()).thenReturn(new PassThroughRequestThrottler(context));

    when(context.getRequestTracker()).thenReturn(new NoopRequestTracker(context));

    when(context.getRequestIdGenerator())
        .thenReturn(Optional.ofNullable(builder.requestIdGenerator));
  }

  public DefaultSession getSession() {
    return session;
  }

  public InternalDriverContext getContext() {
    return context;
  }

  public DriverChannel getChannel(Node node) {
    ChannelPool pool = pools.get(node);
    return pool.next();
  }

  /**
   * Returns the next task that was scheduled on the request handler's admin executor. The test must
   * run it manually.
   */
  public CapturedTimeout nextScheduledTimeout() {
    return timer.getNextTimeout();
  }

  @Override
  public void close() {
    timer.stop();
  }

  public static class Builder {
    private final List<PoolBehavior> poolBehaviors = new ArrayList<>();
    private boolean defaultIdempotence;
    private ProtocolVersion protocolVersion;
    private RequestIdGenerator requestIdGenerator;

    /**
     * Sets the given node as the next one in the query plan; an empty pool will be simulated when
     * it gets used.
     */
    public Builder withEmptyPool(Node node) {
      poolBehaviors.add(new PoolBehavior(node, false));
      return this;
    }

    /**
     * Sets the given node as the next one in the query plan; a channel write failure will be
     * simulated when it gets used.
     */
    public Builder withWriteFailure(Node node, Throwable cause) {
      PoolBehavior behavior = new PoolBehavior(node, true);
      behavior.setWriteFailure(cause);
      poolBehaviors.add(behavior);
      return this;
    }

    /**
     * Sets the given node as the next one in the query plan; the write to the channel will succeed,
     * but a response failure will be simulated immediately after.
     */
    public Builder withResponseFailure(Node node, Throwable cause) {
      PoolBehavior behavior = new PoolBehavior(node, true);
      behavior.setWriteSuccess();
      behavior.setResponseFailure(cause);
      poolBehaviors.add(behavior);
      return this;
    }

    /**
     * Sets the given node as the next one in the query plan; the write to the channel will succeed,
     * and the given response will be simulated immediately after.
     */
    public Builder withResponse(Node node, Frame response) {
      PoolBehavior behavior = new PoolBehavior(node, true);
      behavior.setWriteSuccess();
      behavior.setResponseSuccess(response);
      poolBehaviors.add(behavior);
      return this;
    }

    public Builder withDefaultIdempotence(boolean defaultIdempotence) {
      this.defaultIdempotence = defaultIdempotence;
      return this;
    }

    public Builder withProtocolVersion(ProtocolVersion protocolVersion) {
      this.protocolVersion = protocolVersion;
      return this;
    }

    public Builder withRequestIdGenerator(RequestIdGenerator requestIdGenerator) {
      this.requestIdGenerator = requestIdGenerator;
      return this;
    }

    /**
     * Sets the given node as the next one in the query plan; the test code is responsible of
     * calling the methods on the returned object to complete the write and the query.
     */
    public PoolBehavior customBehavior(Node node) {
      PoolBehavior behavior = new PoolBehavior(node, true);
      poolBehaviors.add(behavior);
      return behavior;
    }

    public RequestHandlerTestHarness build() {
      return new RequestHandlerTestHarness(this);
    }

    private Queue<Node> buildQueryPlan() {
      ConcurrentLinkedQueue<Node> queryPlan = new ConcurrentLinkedQueue<>();
      for (PoolBehavior behavior : poolBehaviors) {
        // We don't want duplicates in the query plan: the only way a node is tried multiple times
        // is if the retry policy returns a RETRY_SAME, the request handler does not re-read from
        // the plan.
        if (!queryPlan.contains(behavior.node)) {
          queryPlan.offer(behavior.node);
        }
      }
      return queryPlan;
    }

    private Map<Node, ChannelPool> buildMockPools() {
      Map<Node, ChannelPool> pools = new ConcurrentHashMap<>();
      Map<Node, OngoingStubbing<DriverChannel>> stubbings = new HashMap<>();
      for (PoolBehavior behavior : poolBehaviors) {
        Node node = behavior.node;
        ChannelPool pool = pools.computeIfAbsent(node, n -> mock(ChannelPool.class));

        // The goal of the code below is to generate the equivalent of:
        //
        //     when(pool.next())
        //       .thenReturn(behavior1.channel)
        //       .thenReturn(behavior2.channel)
        //       ...
        stubbings.compute(
            node,
            (sameNode, previous) -> {
              if (previous == null) {
                previous = when(pool.next());
              }
              return previous.thenReturn(behavior.channel);
            });
      }
      return pools;
    }
  }
}
