/*
 * 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.tests.tez;

import java.io.IOException;
import java.nio.ByteBuffer;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.conf.Configured;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.util.Tool;
import org.apache.hadoop.util.ToolRunner;
import org.apache.hadoop.yarn.api.records.NodeState;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.tez.client.TezClient;
import org.apache.tez.common.Preconditions;
import org.apache.tez.dag.api.*;
import org.apache.tez.dag.api.client.DAGClient;
import org.apache.tez.dag.api.client.DAGStatus;
import org.apache.tez.dag.library.vertexmanager.InputReadyVertexManager;
import org.apache.tez.runtime.api.LogicalOutput;
import org.apache.tez.runtime.api.ObjectRegistry;
import org.apache.tez.runtime.api.ProcessorContext;
import org.apache.tez.runtime.library.api.KeyValueReader;
import org.apache.tez.runtime.library.api.KeyValueWriter;
import org.apache.tez.runtime.library.conf.UnorderedKVEdgeConfig;
import org.apache.tez.runtime.library.output.CelebornUnorderedKVOutput;
import org.apache.tez.runtime.library.output.UnorderedKVOutput;
import org.apache.tez.runtime.library.processor.SimpleProcessor;

public class BroadcastAndOneToOneExample extends Configured implements Tool {
  static String skipLocalityCheck = "-skipLocalityCheck";

  @Override
  public int run(String[] args) throws Exception {
    boolean doLocalityCheck = true;
    if (args.length == 1) {
      if (args[0].equals(skipLocalityCheck)) {
        doLocalityCheck = false;
      } else {
        printUsage();
        throw new TezException("Invalid command line");
      }
    } else if (args.length > 1) {
      printUsage();
      throw new TezException("Invalid command line");
    }

    if (doLocalityCheck
        && getConf()
            .getBoolean(TezConfiguration.TEZ_LOCAL_MODE, TezConfiguration.TEZ_LOCAL_MODE_DEFAULT)) {
      System.out.println("locality check is not valid in local mode. skipping");
      doLocalityCheck = false;
    }

    boolean status = run(getConf(), doLocalityCheck);
    return status ? 0 : 1;
  }

  private static void printUsage() {
    System.err.println("broadcastAndOneToOneExample " + skipLocalityCheck);
    ToolRunner.printGenericCommandUsage(System.err);
  }

  public boolean run(Configuration conf, boolean doLocalityCheck) throws Exception {
    System.out.println("Running BroadcastAndOneToOneExample");
    // conf and UGI
    TezConfiguration tezConf;
    if (conf != null) {
      tezConf = new TezConfiguration(conf);
    } else {
      tezConf = new TezConfiguration();
    }
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    UserGroupInformation.setConfiguration(tezConf);

    // staging dir
    FileSystem fs = FileSystem.get(tezConf);
    String stagingDirStr =
        tezConf.get(
                TezConfiguration.TEZ_AM_STAGING_DIR, TezConfiguration.TEZ_AM_STAGING_DIR_DEFAULT)
            + Path.SEPARATOR
            + "BroadcastAndOneToOneExample"
            + Path.SEPARATOR
            + Long.toString(System.currentTimeMillis());
    Path stagingDir = new Path(stagingDirStr);
    tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, stagingDirStr);
    stagingDir = fs.makeQualified(stagingDir);

    // No need to add jar containing this class as assumed to be part of
    // the tez jars.

    // TEZ-674 Obtain tokens based on the Input / Output paths. For now assuming staging dir
    // is the same filesystem as the one used for Input/Output.
    TezClient tezSession = null;
    // needs session or else TaskScheduler does not hold onto containers
    tezSession = TezClient.create("broadcastAndOneToOneExample", tezConf);
    tezSession.start();

    DAGClient dagClient = null;

    try {
      DAG dag = createDAG(fs, tezConf, stagingDir, doLocalityCheck);

      tezSession.waitTillReady();
      dagClient = tezSession.submitDAG(dag);

      // monitoring
      DAGStatus dagStatus = dagClient.waitForCompletionWithStatusUpdates(null);
      if (dagStatus.getState() != DAGStatus.State.SUCCEEDED) {
        System.out.println("DAG diagnostics: " + dagStatus.getDiagnostics());
        return false;
      }
      return true;
    } finally {
      fs.delete(stagingDir, true);
      tezSession.stop();
    }
  }

  private DAG createDAG(
      FileSystem fs, TezConfiguration tezConf, Path stagingDir, boolean doLocalityCheck)
      throws IOException, YarnException {

    int numBroadcastTasks = 2;
    int numOneToOneTasks = 3;
    if (doLocalityCheck) {
      YarnClient yarnClient = YarnClient.createYarnClient();
      yarnClient.init(tezConf);
      yarnClient.start();
      int numNMs = yarnClient.getNodeReports(NodeState.RUNNING).size();
      yarnClient.stop();
      // create enough 1-1 tasks to run in parallel
      numOneToOneTasks = numNMs - numBroadcastTasks - 1; // 1 AM
      if (numOneToOneTasks < 1) {
        numOneToOneTasks = 1;
      }
    }
    byte[] procByte = {(byte) (doLocalityCheck ? 1 : 0), 1};
    UserPayload procPayload = UserPayload.create(ByteBuffer.wrap(procByte));

    System.out.println("Using " + numOneToOneTasks + " 1-1 tasks");

    Vertex broadcastVertex =
        Vertex.create(
            "Broadcast",
            ProcessorDescriptor.create(InputProcessor.class.getName()),
            numBroadcastTasks);

    Vertex inputVertex =
        Vertex.create(
            "Input",
            ProcessorDescriptor.create(InputProcessor.class.getName()).setUserPayload(procPayload),
            numOneToOneTasks);

    Vertex oneToOneVertex =
        Vertex.create(
            "OneToOne",
            ProcessorDescriptor.create(OneToOneProcessor.class.getName())
                .setUserPayload(procPayload));
    oneToOneVertex.setVertexManagerPlugin(
        VertexManagerPluginDescriptor.create(InputReadyVertexManager.class.getName()));

    UnorderedKVEdgeConfig edgeConf =
        UnorderedKVEdgeConfig.newBuilder(Text.class.getName(), IntWritable.class.getName())
            .setFromConfiguration(tezConf)
            .build();

    DAG dag = DAG.create("BroadcastAndOneToOneExample");
    dag.addVertex(inputVertex)
        .addVertex(broadcastVertex)
        .addVertex(oneToOneVertex)
        .addEdge(
            Edge.create(inputVertex, oneToOneVertex, edgeConf.createDefaultOneToOneEdgeProperty()))
        .addEdge(
            Edge.create(
                broadcastVertex, oneToOneVertex, edgeConf.createDefaultBroadcastEdgeProperty()));
    return dag;
  }

  public static class InputProcessor extends SimpleProcessor {
    Text word = new Text();

    public InputProcessor(ProcessorContext context) {
      super(context);
    }

    @Override
    public void run() throws Exception {
      Preconditions.checkArgument(getOutputs().size() == 1);
      LogicalOutput next = getOutputs().values().iterator().next();
      KeyValueWriter kvWriter;
      if (next instanceof UnorderedKVOutput) {
        kvWriter = ((UnorderedKVOutput) next).getWriter();
      } else {
        kvWriter = ((CelebornUnorderedKVOutput) next).getWriter();
      }
      kvWriter.write(word, new IntWritable(getContext().getTaskIndex()));
      ByteBuffer userPayload =
          getContext().getUserPayload() == null ? null : getContext().getUserPayload().getPayload();
      if (userPayload != null) {
        boolean doLocalityCheck =
            getContext().getUserPayload().getPayload().get(0) > 0 ? true : false;
        if (doLocalityCheck) {
          ObjectRegistry objectRegistry = getContext().getObjectRegistry();
          String entry = String.valueOf(getContext().getTaskIndex());
          objectRegistry.cacheForDAG(entry, entry);
        }
      }
    }
  }

  public static class OneToOneProcessor extends SimpleProcessor {
    Text word = new Text();

    public OneToOneProcessor(ProcessorContext context) {
      super(context);
    }

    @Override
    public void run() throws Exception {
      Preconditions.checkArgument(inputs.size() == 2);

      KeyValueReader inputKvReader = (KeyValueReader) getInputs().get("Input").getReader();
      KeyValueReader broadcastKvReader = (KeyValueReader) getInputs().get("Broadcast").getReader();
      int sum = 0;
      while (broadcastKvReader.next()) {
        sum += ((IntWritable) broadcastKvReader.getCurrentValue()).get();
      }
      while (inputKvReader.next()) {
        sum += ((IntWritable) inputKvReader.getCurrentValue()).get();
      }
      boolean doLocalityCheck =
          getContext().getUserPayload().getPayload().get(0) > 0 ? true : false;
      int broadcastSum = getContext().getUserPayload().getPayload().get(1);
      int expectedSum = broadcastSum + getContext().getTaskIndex();
      System.out.println(
          "Index: "
              + getContext().getTaskIndex()
              + " sum: "
              + sum
              + " expectedSum: "
              + expectedSum
              + " broadcastSum: "
              + broadcastSum);
      Preconditions.checkState((sum == expectedSum), "Sum = " + sum);

      if (doLocalityCheck) {
        ObjectRegistry objectRegistry = getContext().getObjectRegistry();
        String index = (String) objectRegistry.get(String.valueOf(getContext().getTaskIndex()));
        if (index == null || Integer.parseInt(index) != getContext().getTaskIndex()) {
          String msg =
              "Did not find expected local producer "
                  + getContext().getTaskIndex()
                  + " in the same JVM";
          System.out.println(msg);
          throw new TezUncheckedException(msg);
        }
      }
    }
  }
}
