/*
 * 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 org.apache.celeborn.common.protocol;

import static org.junit.Assert.assertEquals;

import org.junit.Test;
import org.roaringbitmap.RoaringBitmap;

public class PartitionLocationSuiteJ {

  private final int partitionId = 0;
  private final int epoch = 0;
  private final String host = "localhost";
  private final int rpcPort = 3;
  private final int pushPort = 1;
  private final int fetchPort = 2;
  private final int replicatePort = 4;
  private final PartitionLocation.Mode mode = PartitionLocation.Mode.PRIMARY;
  private final PartitionLocation peer =
      new PartitionLocation(
          partitionId,
          epoch,
          host,
          rpcPort,
          pushPort,
          fetchPort,
          replicatePort,
          PartitionLocation.Mode.REPLICA);

  @Test
  public void testGetCorrectMode() {
    byte primaryMode = 0;
    byte replicaMode = 1;

    assertEquals(PartitionLocation.Mode.PRIMARY, PartitionLocation.getMode(primaryMode));
    assertEquals(PartitionLocation.Mode.REPLICA, PartitionLocation.getMode(replicaMode));

    for (int i = 2; i < 255; ++i) {
      byte otherMode = (byte) i;
      // Should we return replica mode when the parameter passed in is neither 0 nor 1?
      assertEquals(PartitionLocation.Mode.REPLICA, PartitionLocation.getMode(otherMode));
    }
  }

  @Test
  public void testPartitionIdNotEqualMakePartitionLocationDifferent() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId + 1, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    checkEqual(location1, location2, false);
  }

  @Test
  public void testEpochNotEqualMakePartitionLocationDifferent() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId, epoch + 1, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    checkEqual(location1, location2, false);
  }

  @Test
  public void testHostNotEqualMakePartitionLocationDifferent() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId,
            epoch,
            "remoteHost",
            rpcPort,
            pushPort,
            fetchPort,
            replicatePort,
            mode,
            peer);
    checkEqual(location1, location2, false);
  }

  @Test
  public void testPushPortNotEqualMakePartitionLocationDifferent() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort + 1, fetchPort, replicatePort, mode, peer);
    checkEqual(location1, location2, false);
  }

  @Test
  public void testFetchPortNotEqualMakePartitionLocationDifferent() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort + 1, replicatePort, mode, peer);
    checkEqual(location1, location2, false);
  }

  @Test
  public void testModeNotEqualNeverMakePartitionLocationDifferent() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId,
            epoch,
            host,
            rpcPort,
            pushPort,
            fetchPort,
            replicatePort,
            PartitionLocation.Mode.REPLICA,
            peer);
    PartitionLocation location3 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    checkEqual(location1, location2, true);
    checkEqual(location1, location3, true);
    checkEqual(location2, location3, true);
  }

  @Test
  public void testPeerNotEqualNeverMakePartitionLocationDifferent() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, location1);
    PartitionLocation location3 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    checkEqual(location1, location2, true);
    checkEqual(location1, location3, true);
    checkEqual(location2, location3, true);
  }

  @Test
  public void testAllFieldEqualMakePartitionLocationEqual() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    checkEqual(location1, location2, true);
  }

  @Test
  public void testToStringOutput() {
    PartitionLocation location1 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode);
    PartitionLocation location2 =
        new PartitionLocation(
            partitionId, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode, peer);
    StorageInfo storageInfo =
        new StorageInfo(
            "/mnt/disk/0", StorageInfo.Type.MEMORY, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
    RoaringBitmap bitmap = new RoaringBitmap();
    bitmap.add(1);
    bitmap.add(2);
    bitmap.add(3);

    int partitionId = 1000;
    PartitionLocation location3 =
        new PartitionLocation(
            partitionId,
            epoch,
            host,
            rpcPort,
            pushPort,
            fetchPort,
            replicatePort,
            mode,
            peer,
            storageInfo,
            bitmap);

    String exp1 =
        "PartitionLocation[\n"
            + "  id-epoch:0-0\n"
            + "  host-rpcPort-pushPort-fetchPort-replicatePort:localhost-3-1-2-4\n"
            + "  mode:PRIMARY\n"
            + "  peer:(empty)\n"
            + "  storage hint:StorageInfo{type=MEMORY, mountPoint='', finalResult=false, filePath=null, fileSize=0, chunkOffsets=null}\n"
            + "  mapIdBitMap:{}]";
    String exp2 =
        "PartitionLocation[\n"
            + "  id-epoch:0-0\n"
            + "  host-rpcPort-pushPort-fetchPort-replicatePort:localhost-3-1-2-4\n"
            + "  mode:PRIMARY\n"
            + "  peer:(host-rpcPort-pushPort-fetchPort-replicatePort:localhost-3-1-2-4)\n"
            + "  storage hint:StorageInfo{type=MEMORY, mountPoint='', finalResult=false, filePath=null, fileSize=0, chunkOffsets=null}\n"
            + "  mapIdBitMap:{}]";
    String exp3 =
        "PartitionLocation[\n"
            + "  id-epoch:1000-0\n"
            + "  host-rpcPort-pushPort-fetchPort-replicatePort:localhost-3-1-2-4\n"
            + "  mode:PRIMARY\n"
            + "  peer:(host-rpcPort-pushPort-fetchPort-replicatePort:localhost-3-1-2-4)\n"
            + "  storage hint:StorageInfo{type=MEMORY, mountPoint='/mnt/disk/0', "
            + "finalResult=false, filePath=null, fileSize=0, chunkOffsets=null}\n"
            + "  mapIdBitMap:{1,2,3}]";
    assertEquals(exp1, location1.toString());
    assertEquals(exp2, location2.toString());
    assertEquals(exp3, location3.toString());
  }

  private void checkEqual(
      PartitionLocation location1, PartitionLocation location2, boolean shouldEqual) {
    String errorMessage =
        "Need location1 "
            + location1
            + " and location2 "
            + location2
            + " are "
            + (shouldEqual ? "" : "not ")
            + "equal, but "
            + (shouldEqual ? "not " : "")
            + "equal.";
    assertEquals(errorMessage, shouldEqual, location1.equals(location2));
  }
}
