/*
 * 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.stress.cli;

import alluxio.annotation.SuppressFBWarnings;
import alluxio.client.job.JobGrpcClientUtils;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.job.plan.PlanConfig;
import alluxio.job.wire.JobInfo;
import alluxio.stress.BaseParameters;
import alluxio.stress.StressConstants;
import alluxio.stress.TaskResult;
import alluxio.stress.job.StressBenchConfig;
import alluxio.util.FormatUtils;
import alluxio.util.ShellUtils;

import com.beust.jcommander.JCommander;
import com.beust.jcommander.ParametersDelegate;
import com.fasterxml.jackson.core.JsonParseException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.exc.MismatchedInputException;
import org.HdrHistogram.Histogram;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * Base class for all stress benchmarks.
 *
 * @param <T> the type of task result
 */
public abstract class Benchmark<T extends TaskResult> {
  private static final Logger LOG = LoggerFactory.getLogger(Benchmark.class);

  @ParametersDelegate
  protected BaseParameters mBaseParameters = new BaseParameters();

  /**
   * Get the description of the bench.
   *
   * @return string of the bench description
   */
  public abstract String getBenchDescription();

  /**
   * Runs the test locally, in process.
   *
   * @return the task result
   */
  public abstract T runLocal() throws Exception;

  /**
   * Prepares to run the test.
   */
  // TODO(bowen): When the test runs in cluster mode, the prepare step will execute
  //  both in the command side and on each job worker side. We should separate the logic
  //  into two different calls instead of relying on the same prepare().
  public abstract void prepare() throws Exception;

  /**
   * Perform post-run cleanups.
   */
  public void cleanup() throws Exception {}

  protected static void mainInternal(String[] args, Benchmark benchmark) {
    int exitCode = 0;
    try {
      String result = benchmark.run(args);
      System.out.println(result);
    } catch (Exception e) {
      e.printStackTrace();
      exitCode = -1;
    } finally {
      try {
        benchmark.cleanup();
      } catch (Exception e) {
        e.printStackTrace();
        exitCode = -1;
      }
    }
    System.exit(exitCode);
  }

  /**
   * Generate a {@link StressBenchConfig} as the default JobConfig.
   *
   * @param args arguments
   * @return the JobConfig
   * */
  public PlanConfig generateJobConfig(String[] args) {
    // remove the cluster flag and java opts
    List<String> commandArgs = Arrays.stream(args).filter((s) ->
        !BaseParameters.CLUSTER_FLAG.equals(s) && !s.isEmpty())
        .collect(Collectors.toList());

    commandArgs.addAll(mBaseParameters.mJavaOpts.stream().map(String::trim)
        .collect(Collectors.toList()));
    LOG.info("Generated job config command args: {}", commandArgs);
    String className = this.getClass().getCanonicalName();
    long startDelay = FormatUtils.parseTimeSize(mBaseParameters.mClusterStartDelay);
    return new StressBenchConfig(className, commandArgs, startDelay, mBaseParameters.mClusterLimit);
  }

  /**
   * Runs the test and returns the string output.
   *
   * @param args the command-line args
   * @return the string result output
   */
  public String run(String[] args) throws Exception {
    parseParameters(args);
    return runSingleTask(args);
  }

  protected void parseParameters(String[] args) {
    JCommander jc = new JCommander(this);
    jc.setProgramName(this.getClass().getSimpleName());
    try {
      jc.parse(args);
      if (mBaseParameters.mHelp) {
        System.out.println(getBenchDescription());
        jc.usage();
        System.exit(0);
      }
    } catch (Exception e) {
      LOG.error("Failed to parse command: ", e);
      System.out.println(getBenchDescription());
      jc.usage();
      throw e;
    }
  }

  protected String runSingleTask(String[] args) throws Exception {
    // prepare the benchmark.
    prepare();

    if (!mBaseParameters.mProfileAgent.isEmpty()) {
      mBaseParameters.mJavaOpts.add("-javaagent:" + mBaseParameters.mProfileAgent
          + "=" + BaseParameters.AGENT_OUTPUT_PATH);
    }

    AlluxioConfiguration conf = Configuration.global();
    String className = this.getClass().getCanonicalName();

    if (mBaseParameters.mCluster) {
      // run on job service
      long jobId =
          JobGrpcClientUtils.run(generateJobConfig(args), 0, conf);
      JobInfo jobInfo = JobGrpcClientUtils.getJobStatus(jobId, conf, true);
      String jobResultJson = jobInfo.getResult().toString();
      LOG.info("Collected job results: {}", jobResultJson);
      return jobResultJson;
    }

    // run locally
    if (mBaseParameters.mInProcess) {
      // run in process
      T result = runLocal();
      if (mBaseParameters.mDistributed) {
        String resultsJson = result.toJson();
        LOG.info("This process has results {}", resultsJson);
        return resultsJson;
      }

      // aggregate the results
      String resultAggregated = result.aggregator().aggregate(
          Collections.singletonList(result)).toJson();
      LOG.info("Merged results {}", resultAggregated);
      return resultAggregated;
    } else {
      // Spawn a new process
      List<String> command = new ArrayList<>();
      command.add(conf.get(PropertyKey.HOME) + "/bin/alluxio");
      command.add("runClass");
      command.add(className);
      command.addAll(Arrays.asList(args));
      command.add(BaseParameters.IN_PROCESS_FLAG);
      command.addAll(mBaseParameters.mJavaOpts.stream().map(String::trim)
          .collect(Collectors.toList()));

      LOG.info("running command: " + String.join(" ", command));
      String commandOutput = ShellUtils.execCommand(command.toArray(new String[0]));
      LOG.info("Command output: {}", commandOutput);
      return commandOutput;
    }
  }

  /**
   *
   * @param startMs the start time
   * @param endMs the end time
   * @param nameTransformer function which transforms the type and method into a name. If the
   *                        function returns null, then the method is skipped
   * @return a map of names to statistics
   */
  @SuppressFBWarnings(value = "DMI_HARDCODED_ABSOLUTE_FILENAME")
  protected Map<String, MethodStatistics> processMethodProfiles(long startMs, long endMs,
      Function<ProfileInput, String> nameTransformer) throws IOException {
    Map<String, MethodStatistics> nameStatistics = new HashMap<>();

    try (final BufferedReader reader = new BufferedReader(
        new FileReader(BaseParameters.AGENT_OUTPUT_PATH))) {
      String line;

      long bucketSize = (endMs - startMs) / StressConstants.MAX_TIME_COUNT;

      final ObjectMapper objectMapper = new ObjectMapper();
      while ((line = reader.readLine()) != null) {
        final Map<String, Object> lineMap;
        try {
          lineMap = objectMapper.readValue(line, Map.class);
        } catch (JsonParseException | MismatchedInputException e) {
          // skip the last line of a not completed file
          break;
        }

        final String type = (String) lineMap.get("type");
        final String methodName = (String) lineMap.get("methodName");
        final Number timestampNumber = (Number) lineMap.get("timestamp");
        final Number durationNumber = (Number) lineMap.get("duration");
        final Boolean ttfbFlag = (Boolean) lineMap.get("ttfb");

        if (type == null || methodName == null || timestampNumber == null
            || durationNumber == null || ttfbFlag == null) {
          continue;
        }

        final long timestamp = timestampNumber.longValue();
        final long duration = durationNumber.longValue();
        final boolean ttfb = ttfbFlag;

        if (timestamp <= startMs) {
          continue;
        }

        ProfileInput profileInput = new ProfileInput(type, methodName, ttfb);
        final String name = nameTransformer.apply(profileInput);
        if (name == null) {
          continue;
        }

        if (!nameStatistics.containsKey(name)) {
          nameStatistics.put(name, new MethodStatistics());
        }
        final MethodStatistics statistic = nameStatistics.get(name);

        statistic.mTimeNs.recordValue(duration);
        statistic.mNumSuccess += 1;

        int bucket =
            Math.min(statistic.mMaxTimeNs.length - 1, (int) ((timestamp - startMs) / bucketSize));
        statistic.mMaxTimeNs[bucket] = Math.max(statistic.mMaxTimeNs[bucket], duration);
      }
    }
    return nameStatistics;
  }

  protected static final class ProfileInput {
    private final String mType;
    private final String mMethod;
    private final boolean mIsttfb;

    ProfileInput(String type, String method, boolean isttfb) {
      mType = type;
      mMethod = method;
      mIsttfb = isttfb;
    }

    /**
     * @return class type
     */
    public String getType() {
      return mType;
    }

    /**
     * @return method name
     */
    public String getMethod() {
      return mMethod;
    }

    /**
     * @return is time to first byte
     */
    public boolean getIsttfb() {
      return mIsttfb;
    }
  }

  protected static final class MethodStatistics {
    private final Histogram mTimeNs;
    private int mNumSuccess;
    private final long[] mMaxTimeNs;

    MethodStatistics() {
      mNumSuccess = 0;
      mTimeNs = new Histogram(StressConstants.TIME_HISTOGRAM_MAX,
          StressConstants.TIME_HISTOGRAM_PRECISION);
      mMaxTimeNs = new long[StressConstants.MAX_TIME_COUNT];
      Arrays.fill(mMaxTimeNs, -1);
    }

    /**
     * @return the time histogram
     */
    public Histogram getTimeNs() {
      return mTimeNs;
    }

    /**
     * @return number of successes
     */
    public int getNumSuccess() {
      return mNumSuccess;
    }

    /**
     * @return the max time measurement, over the duration of the test
     */
    public long[] getMaxTimeNs() {
      return mMaxTimeNs;
    }
  }
}
