/*
 * 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 org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyMap;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.loadbalancing.LoadBalancingPolicy;
import com.datastax.oss.driver.api.core.loadbalancing.LoadBalancingPolicy.DistanceReporter;
import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance;
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.context.EventBus;
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.ImmutableList;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSet;
import com.datastax.oss.driver.shaded.guava.common.collect.Lists;
import java.util.Map;
import java.util.Objects;
import java.util.Queue;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

@RunWith(MockitoJUnitRunner.class)
public class LoadBalancingPolicyWrapperTest {

  private DefaultNode node1;
  private DefaultNode node2;
  private DefaultNode node3;

  private Set<DefaultNode> contactPoints;
  private Queue<Node> defaultPolicyQueryPlan;

  @Mock private InternalDriverContext context;
  @Mock private LoadBalancingPolicy policy1;
  @Mock private LoadBalancingPolicy policy2;
  @Mock private LoadBalancingPolicy policy3;
  private EventBus eventBus;
  @Mock private MetadataManager metadataManager;
  @Mock private Metadata metadata;
  @Mock protected MetricsFactory metricsFactory;
  @Captor private ArgumentCaptor<Map<UUID, Node>> initNodesCaptor;

  private LoadBalancingPolicyWrapper wrapper;

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

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

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

    defaultPolicyQueryPlan = Lists.newLinkedList(ImmutableList.of(node3, node2, node1));
    when(policy1.newQueryPlan(null, null)).thenReturn(defaultPolicyQueryPlan);

    eventBus = spy(new EventBus("test"));
    when(context.getEventBus()).thenReturn(eventBus);

    wrapper =
        new LoadBalancingPolicyWrapper(
            context,
            ImmutableMap.of(
                DriverExecutionProfile.DEFAULT_NAME,
                policy1,
                "profile1",
                policy1,
                "profile2",
                policy2,
                "profile3",
                policy3));
  }

  @Test
  public void should_build_query_plan_from_contact_points_before_init() {
    // When
    Queue<Node> queryPlan = wrapper.newQueryPlan();

    // Then
    for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
      verify(policy, never()).newQueryPlan(null, null);
    }
    assertThat(queryPlan).hasSameElementsAs(contactPoints);
  }

  @Test
  public void should_fetch_query_plan_from_policy_after_init() {
    // Given
    wrapper.init();
    for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
      verify(policy).init(anyMap(), any(DistanceReporter.class));
    }

    // When
    Queue<Node> queryPlan = wrapper.newQueryPlan();

    // Then
    // no-arg newQueryPlan() uses the default profile
    verify(policy1).newQueryPlan(null, null);
    assertThat(queryPlan).isEqualTo(defaultPolicyQueryPlan);
  }

  @Test
  public void should_init_policies_with_all_nodes() {
    // Given
    node1.state = NodeState.UP;
    node2.state = NodeState.UNKNOWN;
    node3.state = NodeState.DOWN;

    // When
    wrapper.init();

    // Then
    for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
      verify(policy).init(initNodesCaptor.capture(), any(DistanceReporter.class));
      Map<UUID, Node> initNodes = initNodesCaptor.getValue();
      assertThat(initNodes.values()).containsOnly(node1, node2, node3);
    }
  }

  @Test
  public void should_propagate_distances_from_policies() {
    // Given
    wrapper.init();
    ArgumentCaptor<DistanceReporter> captor1 = ArgumentCaptor.forClass(DistanceReporter.class);
    verify(policy1).init(anyMap(), captor1.capture());
    DistanceReporter distanceReporter1 = captor1.getValue();
    ArgumentCaptor<DistanceReporter> captor2 = ArgumentCaptor.forClass(DistanceReporter.class);
    verify(policy2).init(anyMap(), captor2.capture());
    DistanceReporter distanceReporter2 = captor1.getValue();
    ArgumentCaptor<DistanceReporter> captor3 = ArgumentCaptor.forClass(DistanceReporter.class);
    verify(policy3).init(anyMap(), captor3.capture());
    DistanceReporter distanceReporter3 = captor3.getValue();

    InOrder inOrder = inOrder(eventBus);

    // When
    distanceReporter1.setDistance(node1, NodeDistance.REMOTE);

    // Then
    // first event defines the distance
    inOrder.verify(eventBus).fire(new DistanceEvent(NodeDistance.REMOTE, node1));

    // When
    distanceReporter2.setDistance(node1, NodeDistance.REMOTE);

    // Then
    // event is ignored if the node is already at this distance
    inOrder.verify(eventBus, times(0)).fire(any(DistanceEvent.class));

    // When
    distanceReporter2.setDistance(node1, NodeDistance.LOCAL);

    // Then
    // event is applied if it sets a smaller distance
    inOrder.verify(eventBus).fire(new DistanceEvent(NodeDistance.LOCAL, node1));

    // When
    distanceReporter3.setDistance(node1, NodeDistance.IGNORED);

    // Then
    // event is ignored if the node is already at a closer distance
    inOrder.verify(eventBus, times(0)).fire(any(DistanceEvent.class));
  }

  @Test
  public void should_not_propagate_node_states_to_policies_until_init() {
    // When
    eventBus.fire(NodeStateEvent.changed(NodeState.UNKNOWN, NodeState.UP, node1));

    // Then
    for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
      verify(policy, never()).onUp(node1);
    }
  }

  @Test
  public void should_propagate_node_states_to_policies_after_init() {
    // Given
    wrapper.init();

    // When
    eventBus.fire(NodeStateEvent.changed(NodeState.UNKNOWN, NodeState.UP, node1));

    // Then
    for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
      verify(policy).onUp(node1);
    }
  }

  @Test
  public void should_accumulate_events_during_init_and_replay() throws InterruptedException {
    // Given
    // Hack to obtain concurrency: the main thread releases another thread and blocks; then the
    // other thread fires an event on the bus and unblocks the main thread.
    CountDownLatch eventLatch = new CountDownLatch(1);
    CountDownLatch initLatch = new CountDownLatch(1);

    // When
    Runnable runnable =
        () -> {
          try {
            eventLatch.await();
          } catch (InterruptedException e) {
            throw new RuntimeException(e);
          }
          eventBus.fire(NodeStateEvent.changed(NodeState.UNKNOWN, NodeState.DOWN, node1));
          initLatch.countDown();
        };
    Thread thread = new Thread(runnable);
    thread.start();
    wrapper.init();

    // Then
    // unblock the thread that will fire the event, and waits until it finishes
    eventLatch.countDown();
    boolean ok = initLatch.await(500, TimeUnit.MILLISECONDS);
    assertThat(ok).isTrue();
    for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
      verify(policy).onDown(node1);
    }
    thread.join(500);
    assertThat(thread.isAlive()).isFalse();
  }
}
