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

import static org.mockito.Mockito.mock;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertTrue;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableList.Builder;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Base class for replication strategy tests. Currently only supports testing using the default
 * Murmur3Partitioner.
 */
public class AbstractReplicationStrategyTest {

  private static final Token.Factory partitioner = Token.getFactory("Murmur3Partitioner");

  private static final Cluster.Manager mockManager = mock(Cluster.Manager.class);

  protected static class HostMock extends Host {
    private final InetSocketAddress address;

    private HostMock(InetSocketAddress address, Cluster.Manager manager) {
      super(
          EndPoints.forAddress(address),
          new ConvictionPolicy.DefaultConvictionPolicy.Factory(),
          manager);
      this.address = address;
    }

    private HostMock(InetSocketAddress address, String dc, String rack, Cluster.Manager manager) {
      this(address, manager);
      this.setLocationInfo(dc, rack);
    }

    @Override
    public String toString() {
      return address.toString();
    }

    public InetSocketAddress getMockAddress() {
      return address;
    }

    @Override
    public boolean equals(Object o) {
      if (!(o instanceof HostMock)) return false;

      return address.equals(((HostMock) o).address);
    }

    @Override
    public int hashCode() {
      return address.hashCode();
    }
  }

  protected static Token.Factory partitioner() {
    return partitioner;
  }

  /**
   * Convenience method to quickly create a mock host by a given address. Specified address must be
   * accessible, otherwise a RuntimeException is thrown
   */
  protected static HostMock host(InetSocketAddress address) {
    return new HostMock(address, mockManager);
  }

  /**
   * Convenience method to quickly create a mock host by the given address located in the given
   * datacenter/rack
   */
  protected static HostMock host(InetSocketAddress address, String dc, String rack) {
    return new HostMock(address, dc, rack, mockManager);
  }

  /**
   * Convenience method to cast a Host object into a MockHost. Returns null if parameter host is not
   * a mock
   */
  protected static HostMock asMock(Host host) {
    return (host instanceof HostMock ? (HostMock) host : null);
  }

  /**
   * Convenience method to quickly retrieve a mock host's address as specified if created by the
   * <code>host(...)</code> methods. Returns null if given host is not a mock.
   */
  protected static InetSocketAddress mockAddress(Host host) {
    HostMock mock = asMock(host);
    return mock == null ? null : mock.getMockAddress();
  }

  protected static Token token(String value) {
    return partitioner.fromString(value);
  }

  protected static List<Token> tokens(String... values) {
    Builder<Token> builder = ImmutableList.<Token>builder();
    for (String value : values) {
      builder.add(token(value));
    }
    return builder.build();
  }

  /**
   * Asserts that the replica map for a given token contains the expected list of replica hosts.
   * Hosts are checked in order, replica placement should be an ordered set
   */
  protected static void assertReplicaPlacement(
      Map<Token, Set<Host>> replicaMap, Token token, InetSocketAddress... expected) {
    Set<Host> replicaSet = replicaMap.get(token);
    assertNotNull(replicaSet);
    assertReplicasForToken(replicaSet, expected);
  }

  /** Checks if a given ordered set of replicas matches the expected list of replica hosts */
  protected static void assertReplicasForToken(
      Set<Host> replicaSet, InetSocketAddress... expected) {
    final String message =
        "Contents of replica set: "
            + replicaSet
            + " do not match expected hosts: "
            + Arrays.toString(expected);
    assertEquals(replicaSet.size(), expected.length, message);

    int i = 0;
    for (Host hostReturned : replicaSet) {
      boolean match = true;

      if (!expected[i++].equals(mockAddress(hostReturned))) {
        match = false;
      }
      assertTrue(match, message);
    }
  }

  protected static InetSocketAddress socketAddress(String address) {
    try {
      return new InetSocketAddress(InetAddress.getByName(address), 9042);
    } catch (UnknownHostException ex) {
      throw new RuntimeException(ex);
    }
  }
}
