/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.worker.block;

import static alluxio.worker.block.BlockMasterWorkerServiceTestUtils.createServerWithService;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;

import alluxio.ClientContext;
import alluxio.ConfigurationRule;
import alluxio.Constants;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.FailedToAcquireRegisterLeaseException;
import alluxio.grpc.BlockHeartbeatPRequest;
import alluxio.grpc.BlockHeartbeatPResponse;
import alluxio.grpc.BlockMasterWorkerServiceGrpc;
import alluxio.grpc.BlockStoreLocationProto;
import alluxio.grpc.Command;
import alluxio.grpc.CommandType;
import alluxio.grpc.CommitBlockPRequest;
import alluxio.grpc.CommitBlockPResponse;
import alluxio.grpc.ConfigProperty;
import alluxio.grpc.GetRegisterLeasePRequest;
import alluxio.grpc.GetRegisterLeasePResponse;
import alluxio.grpc.GetWorkerIdPRequest;
import alluxio.grpc.GetWorkerIdPResponse;
import alluxio.grpc.GrpcServer;
import alluxio.grpc.GrpcUtils;
import alluxio.grpc.LocationBlockIdListEntry;
import alluxio.grpc.Metric;
import alluxio.grpc.RegisterWorkerPRequest;
import alluxio.grpc.RegisterWorkerPResponse;
import alluxio.grpc.ServiceType;
import alluxio.grpc.StorageList;
import alluxio.master.MasterClientContext;
import alluxio.retry.CountingRetry;
import alluxio.security.authentication.AuthType;
import alluxio.wire.WorkerNetAddress;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import io.grpc.stub.StreamObserver;
import org.junit.After;
import org.junit.Rule;
import org.junit.Test;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

public class BlockMasterClientTest {

  // test socket address
  private static final InetSocketAddress TEST_SOCKET_ADDRESS =
      new InetSocketAddress("localhost", 9999);

  @Rule
  public ConfigurationRule mConfiguration = new ConfigurationRule(ImmutableMap
      .of(
          PropertyKey.MASTER_RPC_ADDRESSES, ImmutableList.of(TEST_SOCKET_ADDRESS),
          // set retry durations shorter to ensure that tests don't take too long
          PropertyKey.USER_RPC_RETRY_MAX_DURATION, "5s",
          PropertyKey.USER_RPC_RETRY_BASE_SLEEP_MS, "100ms",
          PropertyKey.USER_RPC_RETRY_MAX_SLEEP_MS, "500ms",
          // we don't serve authorization/configuration grpc requests in testing,
          // so make sure clients are not supposed to make those requests
          PropertyKey.SECURITY_AUTHENTICATION_TYPE, AuthType.NOSASL,
          PropertyKey.USER_CONF_CLUSTER_DEFAULT_ENABLED, false),
      Configuration.modifiableGlobal());

  // mock rpc server managed by test suite
  private GrpcServer mServer;

  private final AlluxioConfiguration mConf = Configuration.global();

  @Test
  public void clientInfo() {
    BlockMasterClient client = new BlockMasterClient(
        MasterClientContext.newBuilder(ClientContext.create()).build());

    assertEquals(ServiceType.BLOCK_MASTER_WORKER_SERVICE, client.getRemoteServiceType());
    assertEquals(Constants.BLOCK_MASTER_WORKER_SERVICE_NAME, client.getServiceName());
    assertEquals(Constants.BLOCK_MASTER_WORKER_SERVICE_VERSION, client.getServiceVersion());
  }

  @Test
  public void convertBlockListMapToProtoMergeDirsInSameTier() {
    BlockMasterClient client = new BlockMasterClient(
        MasterClientContext.newBuilder(ClientContext.create()).build());

    Map<BlockStoreLocation, List<Long>> blockMap = new HashMap<>();
    BlockStoreLocation memDir0 = new BlockStoreLocation("MEM", 0);
    blockMap.put(memDir0, Arrays.asList(1L, 2L, 3L));
    BlockStoreLocation memDir1 = new BlockStoreLocation("MEM", 1);
    blockMap.put(memDir1, Arrays.asList(4L, 5L, 6L, 7L));
    BlockStoreLocation ssdDir0 = new BlockStoreLocation("SSD", 0);
    blockMap.put(ssdDir0, Arrays.asList(11L, 12L, 13L, 14L));
    BlockStoreLocation ssdDir1 = new BlockStoreLocation("SSD", 1);
    blockMap.put(ssdDir1, Arrays.asList(15L, 16L, 17L, 18L, 19L));

    // Directories on the same tier will be merged together
    List<LocationBlockIdListEntry> protoList = client.convertBlockListMapToProto(blockMap);
    assertEquals(2, protoList.size());
    BlockStoreLocationProto memLocationProto = BlockStoreLocationProto.newBuilder()
        .setTierAlias("MEM").setMediumType("").build();
    BlockStoreLocationProto ssdLocationProto = BlockStoreLocationProto.newBuilder()
        .setTierAlias("SSD").setMediumType("").build();
    Set<BlockStoreLocationProto> blockLocations = protoList.stream()
        .map(LocationBlockIdListEntry::getKey).collect(Collectors.toSet());
    assertEquals(ImmutableSet.of(memLocationProto, ssdLocationProto), blockLocations);

    LocationBlockIdListEntry firstEntry = protoList.get(0);
    if (firstEntry.getKey().getTierAlias().equals("MEM")) {
      LocationBlockIdListEntry memTierEntry = protoList.get(0);
      List<Long> memProtoBlockList = memTierEntry.getValue().getBlockIdList();
      assertEquals(ImmutableSet.of(1L, 2L, 3L, 4L, 5L, 6L, 7L),
              new HashSet<>(memProtoBlockList));
      LocationBlockIdListEntry ssdTierEntry = protoList.get(1);
      List<Long> ssdProtoBlockList = ssdTierEntry.getValue().getBlockIdList();
      assertEquals(ImmutableSet.of(11L, 12L, 13L, 14L, 15L, 16L, 17L, 18L, 19L),
              new HashSet<>(ssdProtoBlockList));
    } else {
      LocationBlockIdListEntry memTierEntry = protoList.get(1);
      List<Long> memProtoBlockList = memTierEntry.getValue().getBlockIdList();
      assertEquals(ImmutableSet.of(1L, 2L, 3L, 4L, 5L, 6L, 7L),
              new HashSet<>(memProtoBlockList));
      LocationBlockIdListEntry ssdTierEntry = protoList.get(0);
      List<Long> ssdProtoBlockList = ssdTierEntry.getValue().getBlockIdList();
      assertEquals(ImmutableSet.of(11L, 12L, 13L, 14L, 15L, 16L, 17L, 18L, 19L),
              new HashSet<>(ssdProtoBlockList));
    }
  }

  @Test
  public void commitBlock() throws Exception {
    ConcurrentHashMap<Long, Long> committedBlocks = new ConcurrentHashMap<>();
    final long workerId = 1L;
    final long blockId = 2L;
    final long usedBytesOnTier = 1024 * 1024L;
    final long length = 1024 * 1024L;
    final String tierAlias = "MEM";
    final String mediumType = "MEM";

    createMockService(
        new BlockMasterWorkerServiceGrpc.BlockMasterWorkerServiceImplBase() {
          @Override
          public void commitBlock(CommitBlockPRequest request,
                                  StreamObserver<CommitBlockPResponse> responseObserver) {
            long blockId = request.getBlockId();
            long workerId = request.getWorkerId();
            committedBlocks.put(blockId, workerId);
            responseObserver.onNext(CommitBlockPResponse.newBuilder().build());
            responseObserver.onCompleted();
          }
        });

    BlockMasterClient client = new BlockMasterClient(
        MasterClientContext.newBuilder(ClientContext.create(mConf)).build()
    );

    client.commitBlock(workerId, usedBytesOnTier, tierAlias, mediumType, blockId, length);

    assertEquals(1, committedBlocks.size());
    assertEquals(workerId, (long) committedBlocks.get(blockId));
  }

  @Test
  public void getId() throws Exception {
    WorkerNetAddress testExistsAddress = new WorkerNetAddress();
    WorkerNetAddress testNonExistsAddress = new WorkerNetAddress();
    testNonExistsAddress.setHost("1.2.3.4");
    long workerId = 1L;
    Map<WorkerNetAddress, Long> workerIds = ImmutableMap.of(testExistsAddress, workerId);

    createMockService(
        new BlockMasterWorkerServiceGrpc.BlockMasterWorkerServiceImplBase() {
          @Override
          public void getWorkerId(GetWorkerIdPRequest request,
                                  StreamObserver<GetWorkerIdPResponse> responseObserver) {
            WorkerNetAddress address = GrpcUtils.fromProto(request.getWorkerNetAddress());
            Long id = workerIds.get(address);
            if (id == null) {
              id = -1L;
            }
            responseObserver.onNext(GetWorkerIdPResponse.newBuilder().setWorkerId(id).build());
            responseObserver.onCompleted();
          }
        });

    BlockMasterClient client = new BlockMasterClient(
        MasterClientContext.newBuilder(ClientContext.create(mConf)).build()
    );

    assertEquals(workerId, client.getId(testExistsAddress));
    assertEquals(-1L, client.getId(testNonExistsAddress));
  }

  @Test
  public void heartBeat() throws Exception {
    final long workerId = 1L;
    final Map<String, Long> capacityBytesOnTiers = ImmutableMap.of("MEM", 1024 * 1024L);
    final Map<String, Long> usedBytesOnTiers = ImmutableMap.of("MEM", 1024L);
    final List<Long> removedBlocks = ImmutableList.of();
    final Map<BlockStoreLocation, List<Long>> addedBlocks = ImmutableMap.of(
        new BlockStoreLocation("MEM", 0, "MEM"),
        ImmutableList.of(11L, 12L, 13L)
    );
    final Map<String, List<String>> lostStorage = ImmutableMap.of(
        "MEM",
        ImmutableList.of("/tmp/lost")
    );
    final List<Metric> metrics = ImmutableList.of();

    createMockService(
        new BlockMasterWorkerServiceGrpc.BlockMasterWorkerServiceImplBase() {
          @Override
          public void blockHeartbeat(BlockHeartbeatPRequest request,
                                     StreamObserver<BlockHeartbeatPResponse> responseObserver) {

            // verify request data
            assertEquals(workerId, request.getWorkerId());
            assertEquals(usedBytesOnTiers, request.getUsedBytesOnTiersMap());
            assertEquals(removedBlocks, request.getRemovedBlockIdsList());

            // verify added blocks
            for (LocationBlockIdListEntry entry: request.getAddedBlocksList()) {
              BlockStoreLocationProto locationProto = entry.getKey();
              BlockStoreLocation location = new BlockStoreLocation(
                  locationProto.getTierAlias(),
                  0,
                  locationProto.getMediumType()
              );
              List<Long> blocks = addedBlocks.get(location);
              assertTrue(
                  blocks != null && blocks.containsAll(entry.getValue().getBlockIdList())
              );
            }

            // verify lost storage
            assertEquals(lostStorage.entrySet().stream()
                    .collect(Collectors.toMap(Map.Entry::getKey,
                        e -> StorageList.newBuilder().addAllStorage(e.getValue()).build())),
                request.getLostStorageMap());

            responseObserver.onNext(
                BlockHeartbeatPResponse.newBuilder().setCommand(Command.newBuilder()
                .setCommandType(
                    CommandType.Nothing)).build()
            );
            responseObserver.onCompleted();
          }
        });

    BlockMasterClient client = new BlockMasterClient(
        MasterClientContext.newBuilder(ClientContext.create(mConf)).build()
    );

    assertEquals(CommandType.Nothing, client.heartbeat(
        workerId,
        capacityBytesOnTiers,
        usedBytesOnTiers,
        removedBlocks,
        addedBlocks,
        lostStorage,
        metrics).getCommandType());
  }

  @Test
  public void acquireRegisterLeaseFailure() {
    assertThrows(FailedToAcquireRegisterLeaseException.class,
        () -> acquireRegisterLease(false));
  }

  @Test
  public void acquireRegisterLeaseSuccess() throws Exception {
    acquireRegisterLease(true);
  }

  @Test
  public void registerWithoutStream() throws Exception {
    register(false);
  }

  @Test
  public void registerWithStream() throws Exception {
    register(true);
  }

  private void acquireRegisterLease(boolean expectedSuccess) throws Exception {
    createMockService(
        new BlockMasterWorkerServiceGrpc.BlockMasterWorkerServiceImplBase() {
          @Override
          public void requestRegisterLease(
              GetRegisterLeasePRequest request,
              StreamObserver<GetRegisterLeasePResponse> responseObserver) {
            responseObserver.onNext(
                GetRegisterLeasePResponse.newBuilder().setAllowed(expectedSuccess).build());
            responseObserver.onCompleted();
          }
        });

    BlockMasterClient client = new BlockMasterClient(
        MasterClientContext.newBuilder(ClientContext.create(mConf)).build()
    );

    client.acquireRegisterLeaseWithBackoff(
        1L,
        1,
         new CountingRetry(0));
  }

  public void register(boolean stream) throws Exception {
    final long workerId = 1L;
    final List<String> storageTierAliases = ImmutableList.of("MEM");
    final Map<String, Long> totalBytesOnTiers = ImmutableMap.of("MEM", 1024 * 1024L);
    final Map<String, Long> usedBytesOnTiers = ImmutableMap.of("MEM", 1024L);
    // currently have one block at (MEM, 0, MEM)
    // dir-index is set to 0 for recovery from proto
    final Map<BlockStoreLocation, List<Long>> currentBlocksOnLocation = ImmutableMap.of(
        new BlockStoreLocation("MEM", 0, "MEM"), ImmutableList.of(0L)
    );
    final Map<String, List<String>> lostStorage = ImmutableMap.of("MEM", ImmutableList.of());
    final List<ConfigProperty> configList = ImmutableList.of();

    List<Long> registeredWorkerIds = new ArrayList<>();

    createMockService(
        new BlockMasterWorkerServiceGrpc.BlockMasterWorkerServiceImplBase() {
          @Override
          public void registerWorker(RegisterWorkerPRequest request,
                                     StreamObserver<RegisterWorkerPResponse> responseObserver) {

            // verify request data
            assertEquals(workerId, request.getWorkerId());
            assertEquals(storageTierAliases, request.getStorageTiersList());
            assertEquals(totalBytesOnTiers, request.getTotalBytesOnTiersMap());
            assertEquals(usedBytesOnTiers, request.getUsedBytesOnTiersMap());

            // verify that we get all the block information
            for (LocationBlockIdListEntry entry: request.getCurrentBlocksList()) {
              BlockStoreLocationProto locationProto = entry.getKey();
              BlockStoreLocation location = new BlockStoreLocation(
                  locationProto.getTierAlias(), 0, locationProto.getMediumType());
              List<Long> blockIdList = currentBlocksOnLocation.get(location);
              assertTrue(blockIdList != null
                  && blockIdList.containsAll(entry.getValue().getBlockIdList()));
            }

            assertEquals(lostStorage.entrySet().stream()
                .collect(Collectors.toMap(Map.Entry::getKey,
                    e -> StorageList.newBuilder().addAllStorage(e.getValue()).build())),
                request.getLostStorageMap());

            registeredWorkerIds.add(request.getWorkerId());
            responseObserver.onNext(RegisterWorkerPResponse.newBuilder().build());
            responseObserver.onCompleted();
          }

          @Override
          public StreamObserver<RegisterWorkerPRequest> registerWorkerStream(
              StreamObserver<RegisterWorkerPResponse> responseObserver) {
            return new StreamObserver<RegisterWorkerPRequest>() {
              @Override
              public void onNext(RegisterWorkerPRequest value) {
                registeredWorkerIds.add(value.getWorkerId());
                responseObserver.onNext(RegisterWorkerPResponse.newBuilder().build());
              }

              @Override
              public void onError(Throwable t) {
              }

              @Override
              public void onCompleted() {
                responseObserver.onCompleted();
              }
            };
          }
        });

    BlockMasterClient client = new BlockMasterClient(
        MasterClientContext.newBuilder(ClientContext.create(mConf)).build()
    );

    if (stream) {
      client.registerWithStream(
          workerId,
          storageTierAliases,
          totalBytesOnTiers,
          usedBytesOnTiers,
          currentBlocksOnLocation,
          lostStorage,
          configList
      );
    } else {
      client.register(
          workerId,
          storageTierAliases,
          totalBytesOnTiers,
          usedBytesOnTiers,
          currentBlocksOnLocation,
          lostStorage,
          configList);
    }

    assertEquals(1, registeredWorkerIds.size());
    assertEquals(workerId, (long) registeredWorkerIds.get(0));
  }

  @After
  public void cleanUp() {
    if (mServer != null && mServer.isServing()) {
      mServer.shutdown();
    }

    mServer = null;
  }

  // create a mock grpc server that uses delegate to handle rpc calls
  private void createMockService(
      BlockMasterWorkerServiceGrpc.BlockMasterWorkerServiceImplBase delegate) throws IOException {

    // make sure to clean up resources
    cleanUp();

    // set up mock server with custom handler
    mServer = createServerWithService(
        ServiceType.BLOCK_MASTER_WORKER_SERVICE, delegate, TEST_SOCKET_ADDRESS).start();
  }
}
