/*
 * Copyright DataStax, Inc.
 *
 * Licensed 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.dse.driver.internal.core.graph;

import static org.mockito.Mockito.when;

import com.datastax.dse.driver.DseTestFixtures;
import com.datastax.dse.driver.api.core.DseProtocolVersion;
import com.datastax.dse.driver.api.core.config.DseDriverOption;
import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.api.core.Version;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.metadata.Node;
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.context.DefaultDriverContext;
import com.datastax.oss.driver.internal.core.cql.RequestHandlerTestHarness;
import com.datastax.oss.driver.internal.core.servererrors.DefaultWriteTypeRegistry;
import com.datastax.oss.driver.internal.core.session.throttling.PassThroughRequestThrottler;
import com.datastax.oss.driver.internal.core.tracker.NoopRequestTracker;
import com.datastax.oss.protocol.internal.Frame;
import io.netty.channel.EventLoop;
import java.time.Duration;
import javax.annotation.Nullable;
import org.mockito.ArgumentMatchers;
import org.mockito.Mock;

/**
 * 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 GraphRequestHandlerTestHarness extends RequestHandlerTestHarness {

  @Mock DriverExecutionProfile testProfile;

  @Mock DriverExecutionProfile systemQueryExecutionProfile;

  @Mock DefaultDriverContext dseDriverContext;

  @Mock EventLoop eventLoop;

  protected GraphRequestHandlerTestHarness(
      Builder builder,
      @Nullable GraphProtocol graphProtocolForTestConfig,
      Duration graphTimeout,
      @Nullable Version dseVersionForTestMetadata) {
    super(builder);

    // not mocked by RequestHandlerTestHarness, will be used when DseDriverOptions.GRAPH_TIMEOUT
    // is not zero in the config
    when(eventLoopGroup.next()).thenReturn(eventLoop);

    // default graph options as in the reference.conf file
    when(defaultProfile.getString(DseDriverOption.GRAPH_TRAVERSAL_SOURCE, null)).thenReturn("g");
    when(defaultProfile.isDefined(DseDriverOption.GRAPH_SUB_PROTOCOL)).thenReturn(Boolean.FALSE);
    when(defaultProfile.getBoolean(DseDriverOption.GRAPH_IS_SYSTEM_QUERY, false)).thenReturn(false);
    when(defaultProfile.getString(DseDriverOption.GRAPH_NAME, null)).thenReturn("mockGraph");
    when(defaultProfile.getDuration(DseDriverOption.GRAPH_TIMEOUT, Duration.ZERO))
        .thenReturn(graphTimeout);

    when(testProfile.getName()).thenReturn("test-graph");
    when(testProfile.getDuration(DseDriverOption.GRAPH_TIMEOUT, Duration.ZERO))
        .thenReturn(Duration.ofMillis(2L));
    when(testProfile.getString(DefaultDriverOption.REQUEST_CONSISTENCY))
        .thenReturn(DefaultConsistencyLevel.LOCAL_ONE.name());
    when(testProfile.getInt(DefaultDriverOption.REQUEST_PAGE_SIZE)).thenReturn(5000);
    when(testProfile.getString(DefaultDriverOption.REQUEST_SERIAL_CONSISTENCY))
        .thenReturn(DefaultConsistencyLevel.SERIAL.name());
    when(testProfile.getBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE)).thenReturn(false);
    when(testProfile.getBoolean(DefaultDriverOption.PREPARE_ON_ALL_NODES)).thenReturn(true);
    when(testProfile.getString(DseDriverOption.GRAPH_TRAVERSAL_SOURCE, null)).thenReturn("a");
    when(testProfile.isDefined(DseDriverOption.GRAPH_SUB_PROTOCOL))
        .thenReturn(graphProtocolForTestConfig != null);
    // only mock the config if graphProtocolForTestConfig is not null
    if (graphProtocolForTestConfig != null) {
      when(testProfile.getString(DseDriverOption.GRAPH_SUB_PROTOCOL))
          .thenReturn(graphProtocolForTestConfig.toInternalCode());
    }
    when(testProfile.getBoolean(DseDriverOption.GRAPH_IS_SYSTEM_QUERY, false)).thenReturn(false);
    when(testProfile.getString(DseDriverOption.GRAPH_NAME, null)).thenReturn("mockGraph");
    when(testProfile.getString(DseDriverOption.GRAPH_READ_CONSISTENCY_LEVEL, null))
        .thenReturn("LOCAL_TWO");
    when(testProfile.getString(DseDriverOption.GRAPH_WRITE_CONSISTENCY_LEVEL, null))
        .thenReturn("LOCAL_THREE");
    when(config.getProfile("test-graph")).thenReturn(testProfile);

    when(systemQueryExecutionProfile.getName()).thenReturn("graph-system-query");
    when(systemQueryExecutionProfile.getDuration(DseDriverOption.GRAPH_TIMEOUT, Duration.ZERO))
        .thenReturn(Duration.ZERO);
    when(systemQueryExecutionProfile.getString(DefaultDriverOption.REQUEST_CONSISTENCY))
        .thenReturn(DefaultConsistencyLevel.LOCAL_ONE.name());
    when(systemQueryExecutionProfile.getInt(DefaultDriverOption.REQUEST_PAGE_SIZE))
        .thenReturn(5000);
    when(systemQueryExecutionProfile.getString(DefaultDriverOption.REQUEST_SERIAL_CONSISTENCY))
        .thenReturn(DefaultConsistencyLevel.SERIAL.name());
    when(systemQueryExecutionProfile.getBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE))
        .thenReturn(false);
    when(systemQueryExecutionProfile.getBoolean(DefaultDriverOption.PREPARE_ON_ALL_NODES))
        .thenReturn(true);
    when(systemQueryExecutionProfile.getName()).thenReturn("graph-system-query");
    when(systemQueryExecutionProfile.getDuration(DseDriverOption.GRAPH_TIMEOUT, Duration.ZERO))
        .thenReturn(Duration.ofMillis(2));
    when(systemQueryExecutionProfile.getBoolean(DseDriverOption.GRAPH_IS_SYSTEM_QUERY, false))
        .thenReturn(true);
    when(systemQueryExecutionProfile.getString(DseDriverOption.GRAPH_READ_CONSISTENCY_LEVEL, null))
        .thenReturn("LOCAL_TWO");
    when(systemQueryExecutionProfile.getString(DseDriverOption.GRAPH_WRITE_CONSISTENCY_LEVEL, null))
        .thenReturn("LOCAL_THREE");

    when(config.getProfile("graph-system-query")).thenReturn(systemQueryExecutionProfile);

    // need to re-mock everything on the context because the RequestHandlerTestHarness returns a
    // InternalDriverContext and not a DseDriverContext. Couldn't figure out a way with mockito
    // to say "mock this object (this.dseDriverContext), and delegate every call to that
    // other object (this.context), except _this_ call and _this_ and so on"
    // Spy wouldn't work because the spied object has to be of the same type as the final object
    when(dseDriverContext.getConfig()).thenReturn(config);
    when(dseDriverContext.getNettyOptions()).thenReturn(nettyOptions);
    when(dseDriverContext.getLoadBalancingPolicyWrapper()).thenReturn(loadBalancingPolicyWrapper);
    when(dseDriverContext.getRetryPolicy(ArgumentMatchers.anyString())).thenReturn(retryPolicy);
    when(dseDriverContext.getSpeculativeExecutionPolicy(ArgumentMatchers.anyString()))
        .thenReturn(speculativeExecutionPolicy);
    when(dseDriverContext.getCodecRegistry()).thenReturn(CodecRegistry.DEFAULT);
    when(dseDriverContext.getTimestampGenerator()).thenReturn(timestampGenerator);
    when(dseDriverContext.getProtocolVersion()).thenReturn(DseProtocolVersion.DSE_V2);
    when(dseDriverContext.getProtocolVersionRegistry()).thenReturn(protocolVersionRegistry);
    when(dseDriverContext.getConsistencyLevelRegistry())
        .thenReturn(new DefaultConsistencyLevelRegistry());
    when(dseDriverContext.getWriteTypeRegistry()).thenReturn(new DefaultWriteTypeRegistry());
    when(dseDriverContext.getRequestThrottler())
        .thenReturn(new PassThroughRequestThrottler(dseDriverContext));
    when(dseDriverContext.getRequestTracker()).thenReturn(new NoopRequestTracker(dseDriverContext));
    // if DSE Version is specified for test metadata, then we need to mock that up on the context
    if (dseVersionForTestMetadata != null) {
      DseTestFixtures.mockNodesInMetadataWithVersions(
          dseDriverContext, true, dseVersionForTestMetadata);
    }
  }

  @Override
  public DefaultDriverContext getContext() {
    return dseDriverContext;
  }

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

  public static class Builder extends RequestHandlerTestHarness.Builder {

    private GraphProtocol graphProtocolForTestConfig;
    private Duration graphTimeout = Duration.ZERO;
    private Version dseVersionForTestMetadata;

    public GraphRequestHandlerTestHarness.Builder withGraphProtocolForTestConfig(
        GraphProtocol protocol) {
      this.graphProtocolForTestConfig = protocol;
      return this;
    }

    public GraphRequestHandlerTestHarness.Builder withDseVersionInMetadata(Version dseVersion) {
      this.dseVersionForTestMetadata = dseVersion;
      return this;
    }

    public GraphRequestHandlerTestHarness.Builder withGraphTimeout(Duration globalTimeout) {
      this.graphTimeout = globalTimeout;
      return this;
    }

    @Override
    public GraphRequestHandlerTestHarness.Builder withEmptyPool(Node node) {
      super.withEmptyPool(node);
      return this;
    }

    @Override
    public GraphRequestHandlerTestHarness.Builder withWriteFailure(Node node, Throwable cause) {
      super.withWriteFailure(node, cause);
      return this;
    }

    @Override
    public GraphRequestHandlerTestHarness.Builder withResponseFailure(Node node, Throwable cause) {
      super.withResponseFailure(node, cause);
      return this;
    }

    @Override
    public GraphRequestHandlerTestHarness.Builder withResponse(Node node, Frame response) {
      super.withResponse(node, response);
      return this;
    }

    @Override
    public GraphRequestHandlerTestHarness.Builder withDefaultIdempotence(
        boolean defaultIdempotence) {
      super.withDefaultIdempotence(defaultIdempotence);
      return this;
    }

    @Override
    public GraphRequestHandlerTestHarness.Builder withProtocolVersion(
        ProtocolVersion protocolVersion) {
      super.withProtocolVersion(protocolVersion);
      return this;
    }

    @Override
    public GraphRequestHandlerTestHarness build() {
      return new GraphRequestHandlerTestHarness(
          this, graphProtocolForTestConfig, graphTimeout, dseVersionForTestMetadata);
    }
  }
}
