/*
 * 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.oss.driver.internal.core.metadata;

import static com.datastax.oss.driver.Assertions.assertThat;
import static com.datastax.oss.driver.Assertions.assertThatStage;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

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.metadata.Metadata;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metadata.NodeState;
import com.datastax.oss.driver.internal.core.adminrequest.AdminResult;
import com.datastax.oss.driver.internal.core.adminrequest.AdminRow;
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.metrics.MetricsFactory;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
import com.datastax.oss.driver.shaded.guava.common.collect.Iterators;
import com.tngtech.java.junit.dataprovider.DataProvider;
import com.tngtech.java.junit.dataprovider.DataProviderRunner;
import com.tngtech.java.junit.dataprovider.UseDataProvider;
import io.netty.channel.EventLoop;
import java.time.Duration;
import java.util.ArrayDeque;
import java.util.Arrays;
import java.util.Map;
import java.util.Objects;
import java.util.Queue;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

@RunWith(DataProviderRunner.class)
public class SchemaAgreementCheckerTest {

  private static final UUID VERSION1 = UUID.randomUUID();
  private static final UUID VERSION2 = UUID.randomUUID();

  private static final UUID NODE_2_HOST_ID = UUID.randomUUID();

  @Mock private InternalDriverContext context;
  @Mock private DriverConfig config;
  @Mock private DriverExecutionProfile defaultConfig;
  @Mock private DriverChannel channel;
  @Mock private EventLoop eventLoop;
  @Mock private MetadataManager metadataManager;
  @Mock private MetricsFactory metricsFactory;
  @Mock private Metadata metadata;
  @Mock private DefaultNode node1;
  @Mock private DefaultNode node2;

  @Before
  public void setup() {
    MockitoAnnotations.initMocks(this);
    when(context.getMetricsFactory()).thenReturn(metricsFactory);

    node1 = TestNodeFactory.newNode(1, context);
    node2 = TestNodeFactory.newNode(2, NODE_2_HOST_ID, context);

    when(defaultConfig.getDuration(DefaultDriverOption.CONTROL_CONNECTION_TIMEOUT))
        .thenReturn(Duration.ofSeconds(1));
    when(defaultConfig.getDuration(DefaultDriverOption.CONTROL_CONNECTION_AGREEMENT_INTERVAL))
        .thenReturn(Duration.ofMillis(200));
    when(defaultConfig.getDuration(DefaultDriverOption.CONTROL_CONNECTION_AGREEMENT_TIMEOUT))
        .thenReturn(Duration.ofSeconds(10));
    when(defaultConfig.getBoolean(DefaultDriverOption.CONTROL_CONNECTION_AGREEMENT_WARN))
        .thenReturn(true);
    when(config.getDefaultProfile()).thenReturn(defaultConfig);
    when(context.getConfig()).thenReturn(config);

    Map<UUID, Node> nodes =
        ImmutableMap.of(
            Objects.requireNonNull(node1.getHostId()),
            node1,
            Objects.requireNonNull(node2.getHostId()),
            node2);
    when(metadata.getNodes()).thenReturn(nodes);
    when(metadataManager.getMetadata()).thenReturn(metadata);
    when(context.getMetadataManager()).thenReturn(metadataManager);

    node2.state = NodeState.UP;

    when(eventLoop.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)))
        .thenAnswer(
            invocation -> { // Ignore delay and run immediately:
              Runnable task = invocation.getArgument(0);
              task.run();
              return null;
            });
    when(channel.eventLoop()).thenReturn(eventLoop);
  }

  @Test
  public void should_skip_if_timeout_is_zero() {
    // Given
    when(defaultConfig.getDuration(DefaultDriverOption.CONTROL_CONNECTION_AGREEMENT_TIMEOUT))
        .thenReturn(Duration.ZERO);
    TestSchemaAgreementChecker checker = new TestSchemaAgreementChecker(channel, context);

    // When
    CompletionStage<Boolean> future = checker.run();

    // Then
    assertThatStage(future).isSuccess(b -> assertThat(b).isFalse());
  }

  @Test
  public void should_succeed_if_only_one_node() {
    // Given
    TestSchemaAgreementChecker checker = new TestSchemaAgreementChecker(channel, context);
    checker.stubQueries(
        new StubbedQuery(
            "SELECT schema_version FROM system.local WHERE key='local'",
            mockResult(mockLocalRow(VERSION1))),
        new StubbedQuery("SELECT * FROM system.peers", mockResult(/*empty*/ )));

    // When
    CompletionStage<Boolean> future = checker.run();

    // Then
    assertThatStage(future).isSuccess(b -> assertThat(b).isTrue());
  }

  @Test
  public void should_succeed_if_versions_match_on_first_try() {
    // Given
    TestSchemaAgreementChecker checker = new TestSchemaAgreementChecker(channel, context);
    checker.stubQueries(
        new StubbedQuery(
            "SELECT schema_version FROM system.local WHERE key='local'",
            mockResult(mockLocalRow(VERSION1))),
        new StubbedQuery("SELECT * FROM system.peers", mockResult(mockValidPeerRow(VERSION1))));

    // When
    CompletionStage<Boolean> future = checker.run();

    // Then
    assertThatStage(future).isSuccess(b -> assertThat(b).isTrue());
  }

  @Test
  public void should_ignore_down_peers() {
    // Given
    TestSchemaAgreementChecker checker = new TestSchemaAgreementChecker(channel, context);
    node2.state = NodeState.DOWN;
    checker.stubQueries(
        new StubbedQuery(
            "SELECT schema_version FROM system.local WHERE key='local'",
            mockResult(mockLocalRow(VERSION1))),
        new StubbedQuery("SELECT * FROM system.peers", mockResult(mockValidPeerRow(VERSION2))));

    // When
    CompletionStage<Boolean> future = checker.run();

    // Then
    assertThatStage(future).isSuccess(b -> assertThat(b).isTrue());
  }

  @DataProvider
  public static Object[][] malformedPeer() {
    return new Object[][] {
      // missing host id
      {mockPeerRow(null, VERSION2, true, true, true, true)},
      // missing schema version
      {mockPeerRow(NODE_2_HOST_ID, null, true, true, true, true)},
      // missing datacenter
      {mockPeerRow(NODE_2_HOST_ID, VERSION2, false, true, true, true)},
      // missing rack
      {mockPeerRow(NODE_2_HOST_ID, VERSION2, true, false, true, true)},
      // missing RPC address
      {mockPeerRow(NODE_2_HOST_ID, VERSION2, true, true, false, true)},
      // missing tokens
      {mockPeerRow(NODE_2_HOST_ID, VERSION2, true, true, true, false)},
    };
  }

  @Test
  @UseDataProvider("malformedPeer")
  public void should_ignore_malformed_rows(AdminRow malformedPeer) {
    // Given
    TestSchemaAgreementChecker checker = new TestSchemaAgreementChecker(channel, context);
    checker.stubQueries(
        new StubbedQuery(
            "SELECT schema_version FROM system.local WHERE key='local'",
            mockResult(mockLocalRow(VERSION1))),
        new StubbedQuery("SELECT * FROM system.peers", mockResult(malformedPeer)));

    // When
    CompletionStage<Boolean> future = checker.run();

    // Then
    assertThatStage(future).isSuccess(b -> assertThat(b).isTrue());
  }

  @Test
  public void should_reschedule_if_versions_do_not_match_on_first_try() {
    // Given
    TestSchemaAgreementChecker checker = new TestSchemaAgreementChecker(channel, context);
    checker.stubQueries(
        // First round
        new StubbedQuery(
            "SELECT schema_version FROM system.local WHERE key='local'",
            mockResult(mockLocalRow(VERSION1))),
        new StubbedQuery("SELECT * FROM system.peers", mockResult(mockValidPeerRow(VERSION2))),

        // Second round
        new StubbedQuery(
            "SELECT schema_version FROM system.local WHERE key='local'",
            mockResult(mockLocalRow(VERSION1))),
        new StubbedQuery("SELECT * FROM system.peers", mockResult(mockValidPeerRow(VERSION1))));

    // When
    CompletionStage<Boolean> future = checker.run();

    // Then
    assertThatStage(future).isSuccess(b -> assertThat(b).isTrue());
  }

  @Test
  public void should_fail_if_versions_do_not_match_after_timeout() {
    // Given
    when(defaultConfig.getDuration(DefaultDriverOption.CONTROL_CONNECTION_AGREEMENT_TIMEOUT))
        .thenReturn(Duration.ofNanos(10));
    TestSchemaAgreementChecker checker = new TestSchemaAgreementChecker(channel, context);
    checker.stubQueries(
        new StubbedQuery(
            "SELECT schema_version FROM system.local WHERE key='local'",
            mockResult(mockLocalRow(VERSION1))),
        new StubbedQuery("SELECT * FROM system.peers", mockResult(mockValidPeerRow(VERSION1))));

    // When
    CompletionStage<Boolean> future = checker.run();

    // Then
    assertThatStage(future).isSuccess(b -> assertThat(b).isFalse());
  }

  /** Extend to mock the query execution logic. */
  private static class TestSchemaAgreementChecker extends SchemaAgreementChecker {

    private final Queue<StubbedQuery> queries = new ArrayDeque<>();

    TestSchemaAgreementChecker(DriverChannel channel, InternalDriverContext context) {
      super(channel, context, "test");
    }

    private void stubQueries(StubbedQuery... queries) {
      this.queries.addAll(Arrays.asList(queries));
    }

    @Override
    protected CompletionStage<AdminResult> query(String queryString) {
      StubbedQuery nextQuery = queries.poll();
      assertThat(nextQuery).isNotNull();
      assertThat(queryString).isEqualTo(nextQuery.queryString);
      return CompletableFuture.completedFuture(nextQuery.result);
    }
  }

  private static class StubbedQuery {
    private final String queryString;
    private final AdminResult result;

    private StubbedQuery(String queryString, AdminResult result) {
      this.queryString = queryString;
      this.result = result;
    }
  }

  private AdminRow mockLocalRow(@SuppressWarnings("SameParameterValue") UUID schemaVersion) {
    AdminRow row = mock(AdminRow.class);
    when(row.getUuid("host_id")).thenReturn(node1.getHostId());
    when(row.getUuid("schema_version")).thenReturn(schemaVersion);
    return row;
  }

  private AdminRow mockValidPeerRow(UUID schemaVersion) {
    return mockPeerRow(node2.getHostId(), schemaVersion, true, true, true, true);
  }

  private static AdminRow mockPeerRow(
      UUID hostId,
      UUID schemaVersion,
      boolean hasDatacenter,
      boolean hasRack,
      boolean hasRpcAddress,
      boolean hasTokens) {
    AdminRow row = mock(AdminRow.class);
    when(row.getUuid("host_id")).thenReturn(hostId);
    when(row.isNull("host_id")).thenReturn(hostId == null);
    when(row.getUuid("schema_version")).thenReturn(schemaVersion);
    when(row.isNull("schema_version")).thenReturn(schemaVersion == null);
    when(row.isNull("data_center")).thenReturn(!hasDatacenter);
    when(row.isNull("rack")).thenReturn(!hasRack);
    when(row.isNull("tokens")).thenReturn(!hasTokens);
    when(row.isNull("rpc_address")).thenReturn(!hasRpcAddress);
    when(row.isNull("native_address")).thenReturn(true);
    when(row.isNull("native_port")).thenReturn(true);
    return row;
  }

  private AdminResult mockResult(AdminRow... rows) {
    AdminResult result = mock(AdminResult.class);
    when(result.iterator()).thenReturn(Iterators.forArray(rows));
    return result;
  }
}
