/*
 * 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.plugin.flink.tiered;

import java.net.UnknownHostException;
import java.util.Collection;
import java.util.concurrent.CompletableFuture;

import org.apache.flink.api.common.JobID;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
import org.apache.flink.runtime.executiongraph.ExecutionGraphID;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleDescriptor;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleHandler;
import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.util.Utils$;
import org.apache.celeborn.plugin.flink.ShuffleResource;
import org.apache.celeborn.plugin.flink.ShuffleResourceDescriptor;
import org.apache.celeborn.plugin.flink.utils.FlinkUtils;

public class CelebornTierMasterAgentSuiteJ {
  private static final Logger LOG = LoggerFactory.getLogger(CelebornTierMasterAgentSuiteJ.class);
  private CelebornTierMasterAgent masterAgent;

  @Before
  public void setUp() {
    Configuration configuration = new Configuration();
    int startPort = Utils$.MODULE$.selectRandomInt(1024, 65535);
    configuration.setInteger("celeborn.master.port", startPort);
    configuration.setString("celeborn.master.endpoints", "localhost:" + startPort);
    masterAgent = createMasterAgent(configuration);
  }

  @Test
  public void testRegisterJob() {
    TierShuffleHandler tierShuffleHandler = createTierShuffleHandler();
    JobID jobID = JobID.generate();
    masterAgent.registerJob(jobID, tierShuffleHandler);

    // reRunRegister job
    try {
      masterAgent.registerJob(jobID, tierShuffleHandler);
      Assert.fail("should throw exception if double register job");
    } catch (Exception e) {
      Assert.assertTrue(true);
    }

    // unRegister job
    masterAgent.unregisterJob(jobID);
    masterAgent.registerJob(jobID, tierShuffleHandler);
  }

  private static TierShuffleHandler createTierShuffleHandler() {
    return new TierShuffleHandler() {

      @Override
      public CompletableFuture<?> onReleasePartitions(
          Collection<TieredStoragePartitionId> collection) {
        return CompletableFuture.completedFuture(null);
      }

      @Override
      public void onFatalError(Throwable throwable) {
        System.exit(-1);
      }
    };
  }

  @Test
  public void testRegisterPartitionWithProducer() {
    JobID jobID = JobID.generate();
    TierShuffleHandler tierShuffleHandler = createTierShuffleHandler();
    masterAgent.registerJob(jobID, tierShuffleHandler);

    ExecutionAttemptID executionAttemptID =
        new ExecutionAttemptID(
            new ExecutionGraphID(), new ExecutionVertexID(new JobVertexID(0L, 0L), 0), 0);
    ResultPartitionID resultPartitionID =
        new ResultPartitionID(
            new IntermediateResultPartitionID(new IntermediateDataSetID(), 0), executionAttemptID);
    TierShuffleDescriptor tierShuffleDescriptor =
        masterAgent.addPartitionAndGetShuffleDescriptor(jobID, resultPartitionID);
    Assert.assertTrue(tierShuffleDescriptor instanceof TierShuffleDescriptorImpl);
    ShuffleResource shuffleResource =
        ((TierShuffleDescriptorImpl) tierShuffleDescriptor).getShuffleResource();
    ShuffleResourceDescriptor mapPartitionShuffleDescriptor =
        shuffleResource.getMapPartitionShuffleDescriptor();

    Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
    Assert.assertEquals(0, mapPartitionShuffleDescriptor.getPartitionId());
    Assert.assertEquals(0, mapPartitionShuffleDescriptor.getAttemptId());
    Assert.assertEquals(0, mapPartitionShuffleDescriptor.getMapId());

    // use same partition id
    tierShuffleDescriptor =
        masterAgent.addPartitionAndGetShuffleDescriptor(jobID, resultPartitionID);
    mapPartitionShuffleDescriptor =
        ((TierShuffleDescriptorImpl) tierShuffleDescriptor)
            .getShuffleResource()
            .getMapPartitionShuffleDescriptor();
    Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
    Assert.assertEquals(0, mapPartitionShuffleDescriptor.getMapId());
    Assert.assertEquals(1, mapPartitionShuffleDescriptor.getPartitionId());
    Assert.assertEquals(1, mapPartitionShuffleDescriptor.getAttemptId());

    // use another partition number
    tierShuffleDescriptor =
        masterAgent.addPartitionAndGetShuffleDescriptor(
            jobID,
            new ResultPartitionID(
                new IntermediateResultPartitionID(
                    resultPartitionID.getPartitionId().getIntermediateDataSetID(), 1),
                executionAttemptID));
    mapPartitionShuffleDescriptor =
        ((TierShuffleDescriptorImpl) tierShuffleDescriptor)
            .getShuffleResource()
            .getMapPartitionShuffleDescriptor();
    Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
    Assert.assertEquals(1, mapPartitionShuffleDescriptor.getMapId());
    Assert.assertEquals(2, mapPartitionShuffleDescriptor.getPartitionId());
    Assert.assertEquals(0, mapPartitionShuffleDescriptor.getAttemptId());
  }

  @Test
  public void testRegisterMultipleJobs() throws UnknownHostException {
    JobID jobID1 = JobID.generate();
    TierShuffleHandler tierShuffleHandler1 = createTierShuffleHandler();
    masterAgent.registerJob(jobID1, tierShuffleHandler1);

    JobID jobID2 = JobID.generate();
    TierShuffleHandler tierShuffleHandler2 = createTierShuffleHandler();
    masterAgent.registerJob(jobID2, tierShuffleHandler2);

    IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
    ResultPartitionID resultPartitionID = new ResultPartitionID();
    TierShuffleDescriptor tierShuffleDescriptor1 =
        masterAgent.addPartitionAndGetShuffleDescriptor(jobID1, resultPartitionID);

    // use same partition id but different jobId
    TierShuffleDescriptor tierShuffleDescriptor2 =
        masterAgent.addPartitionAndGetShuffleDescriptor(jobID2, resultPartitionID);

    Assert.assertEquals(
        ((TierShuffleDescriptorImpl) tierShuffleDescriptor1)
            .getShuffleResource()
            .getMapPartitionShuffleDescriptor()
            .getShuffleId(),
        0);
    Assert.assertEquals(
        ((TierShuffleDescriptorImpl) tierShuffleDescriptor2)
            .getShuffleResource()
            .getMapPartitionShuffleDescriptor()
            .getShuffleId(),
        1);
  }

  @After
  public void tearDown() {
    if (masterAgent != null) {
      try {
        masterAgent.close();
      } catch (Exception e) {
        LOG.warn(e.getMessage(), e);
      }
    }
  }

  public CelebornTierMasterAgent createMasterAgent(Configuration configuration) {
    CelebornConf conf = FlinkUtils.toCelebornConf(configuration);
    return new CelebornTierMasterAgent(conf);
  }
}
