/*
 * 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.conf.Configuration;
import alluxio.stress.worker.IOTaskResult;
import alluxio.stress.worker.UfsIOParameters;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.UnderFileSystemConfiguration;
import alluxio.util.CommonUtils;
import alluxio.util.ExceptionUtils;
import alluxio.util.FormatUtils;
import alluxio.util.executor.ExecutorServiceFactories;
import alluxio.util.io.PathUtils;

import com.beust.jcommander.ParametersDelegate;
import com.google.common.collect.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * A benchmark tool measuring the IO to UFS.
 * */
public class UfsIOBench extends Benchmark<IOTaskResult> {
  private static final Logger LOG = LoggerFactory.getLogger(UfsIOBench.class);
  private static final int BUFFER_SIZE = 1024 * 1024;
  private static final String TEST_DIR_NAME = "UfsIOTest";

  @ParametersDelegate
  private final UfsIOParameters mParameters = new UfsIOParameters();

  private final UUID mTaskId = UUID.randomUUID();
  private String mDataDir;

  @Override
  public String getBenchDescription() {
    return String.join("\n", ImmutableList.of(
        "A benchmarking tool for the I/O between Alluxio and UFS.",
        "This test will measure the I/O throughput between Alluxio workers and "
            + "the specified UFS path. Each worker will create concurrent clients to "
            + "first generate test files of the specified size then read those files. "
            + "The write/read I/O throughput will be measured in the process.",
        "",
        "Example:",
        "# This invokes the I/O benchmark to HDFS in the Alluxio cluster",
        "# 2 workers will be used",
        "# 2 concurrent clients will be created on each worker",
        "# Each thread is writing then reading 512m of data",
        "$ bin/alluxio runUfsIOTest --path hdfs://<hdfs-address> --cluster --cluster-limit 2 \\",
        " --io-size 512m --threads 2",
        ""
    ));
  }

  @Override
  public IOTaskResult runLocal() throws Exception {
    // UfsIOBench is invoked from the job worker then runs in a standalone process
    // The process type should be set to keep consistent
    boolean switched = CommonUtils.PROCESS_TYPE.compareAndSet(CommonUtils.ProcessType.CLIENT,
        CommonUtils.ProcessType.JOB_WORKER);

    LOG.debug("Running locally with {} threads", mParameters.mThreads);
    ExecutorService pool = null;
    IOTaskResult result = null;
    try {
      pool = ExecutorServiceFactories.fixedThreadPool("bench-io-thread", mParameters.mThreads)
                      .create();

      result = runIOBench(pool);
      LOG.debug("IO benchmark finished with result: {}", result);
      // Aggregate the task results
      return result;
    } catch (Exception e) {
      if (result == null) {
        LOG.error("Failed run UFS IO benchmark on path {}", mParameters.mPath, e);
        result = new IOTaskResult();
        result.setParameters(mParameters);
        result.setBaseParameters(mBaseParameters);
        result.addError(ExceptionUtils.asPlainText(e));
      }
      return result;
    } finally {
      if (pool != null) {
        pool.shutdownNow();
        pool.awaitTermination(30, TimeUnit.SECONDS);
      }
      if (switched) {
        // restore the process type if it was switched
        CommonUtils.PROCESS_TYPE.set(CommonUtils.ProcessType.CLIENT);
      }
    }
  }

  @Override
  public void prepare() {
    if (mParameters.mUseUfsConf && !mBaseParameters.mCluster) {
      throw new IllegalArgumentException(String.format(
          "%s can not use the ufs conf if it is not running in cluster mode",
          getClass().getName()));
    }
    mDataDir = PathUtils.concatPath(mParameters.mPath, TEST_DIR_NAME);
  }

  /**
   * @param args command-line arguments
   */
  public static void main(String[] args) {
    mainInternal(args, new UfsIOBench());
  }

  private String getFilePath(int idx) {
    return PathUtils.concatPath(mDataDir, String.format("io-benchmark-%s-%d", mTaskId, idx));
  }

  private IOTaskResult runIOBench(ExecutorService pool) throws Exception {
    IOTaskResult writeTaskResult = write(pool);
    if (writeTaskResult.getPoints().size() == 0) {
      LOG.error("Failed to write any files. Abort the test.");
      return writeTaskResult;
    }
    IOTaskResult readTaskResult = read(pool);
    cleanUp();
    return writeTaskResult.merge(readTaskResult);
  }

  private void cleanUp() throws IOException {
    UnderFileSystemConfiguration ufsConf = UnderFileSystemConfiguration
        .defaults(Configuration.global())
        .createMountSpecificConf(mParameters.mConf);
    UnderFileSystem ufs = UnderFileSystem.Factory.create(mDataDir, ufsConf);

    for (int i = 0; i < mParameters.mThreads; i++) {
      ufs.deleteFile(getFilePath(i));
    }
  }

  private IOTaskResult read(ExecutorService pool)
          throws InterruptedException, ExecutionException {
    UnderFileSystemConfiguration ufsConf;
    UnderFileSystem ufs;
    int numThreads;
    long ioSizeBytes;
    try {
      // Use multiple threads to saturate the bandwidth of this worker
      numThreads = mParameters.mThreads;
      ioSizeBytes = FormatUtils.parseSpaceSize(mParameters.mDataSize);
      ufsConf = UnderFileSystemConfiguration.defaults(Configuration.global())
              .createMountSpecificConf(mParameters.mConf);
      ufs = UnderFileSystem.Factory.create(mDataDir, ufsConf);
      if (!ufs.exists(mDataDir)) {
        // If the directory does not exist, there's no point proceeding
        throw new IOException(String.format("The target directory %s does not exist!", mDataDir));
      }
    } catch (Exception e) {
      LOG.error("Failed to access UFS path {}", mDataDir);
      // If the UFS path is not valid, abort the test
      IOTaskResult result = new IOTaskResult();
      result.setParameters(mParameters);
      result.setBaseParameters(mBaseParameters);
      result.addError(ExceptionUtils.asPlainText(e));
      return result;
    }

    List<CompletableFuture<IOTaskResult>> futures = new ArrayList<>();
    for (int i = 0; i < numThreads; i++) {
      final int idx = i;
      CompletableFuture<IOTaskResult> future = CompletableFuture.supplyAsync(() -> {
        IOTaskResult result = new IOTaskResult();
        result.setBaseParameters(mBaseParameters);
        result.setParameters(mParameters);
        long startTime = CommonUtils.getCurrentMs();

        String filePath = getFilePath(idx);
        LOG.debug("Reading filePath={}", filePath);

        long readBytes = 0;
        InputStream inStream = null;
        try {
          inStream = ufs.open(filePath);
          byte[] buf = new byte[BUFFER_SIZE];
          int readBufBytes;
          while (readBytes < ioSizeBytes && (readBufBytes = inStream.read(buf)) > 0) {
            readBytes += readBufBytes;
          }

          long endTime = CommonUtils.getCurrentMs();
          double duration = (endTime - startTime) / 1000.0; // convert to second
          IOTaskResult.Point p = new IOTaskResult.Point(IOTaskResult.IOMode.READ,
                  duration, readBytes);
          result.addPoint(p);
          LOG.debug("Read task finished {}", p);
        } catch (Exception e) {
          LOG.error("Failed to read {}", filePath, e);
          result.addError(ExceptionUtils.asPlainText(e));
        } finally {
          if (inStream != null) {
            try {
              inStream.close();
            } catch (IOException e) {
              LOG.warn("Failed to close read stream {}", filePath, e);
              result.addError(e.getMessage());
            }
          }
        }

        return result;
      }, pool);
      futures.add(future);
    }

    // Collect the result
    CompletableFuture[] cfs = futures.toArray(new CompletableFuture[0]);
    List<IOTaskResult> results = CompletableFuture.allOf(cfs)
            .thenApply(f -> futures.stream()
                    .map(CompletableFuture::join)
                    .collect(Collectors.toList())
            ).get();

    return IOTaskResult.reduceList(results);
  }

  private IOTaskResult write(ExecutorService pool)
          throws InterruptedException, ExecutionException {
    UnderFileSystemConfiguration ufsConf;
    UnderFileSystem ufs;
    int numThreads;
    long ioSizeBytes;
    try {
      // Use multiple threads to saturate the bandwidth of this worker
      numThreads = mParameters.mThreads;
      ioSizeBytes = FormatUtils.parseSpaceSize(mParameters.mDataSize);
      ufsConf = UnderFileSystemConfiguration.defaults(Configuration.global())
              .createMountSpecificConf(mParameters.mConf);
      // Create a subdir for the IO
      ufs = UnderFileSystem.Factory.create(mDataDir, ufsConf);
      if (!ufs.exists(mDataDir)) {
        LOG.debug("Prepare directory {}", mDataDir);
        ufs.mkdirs(mDataDir);
      }
    } catch (Exception e) {
      LOG.error("Failed to prepare base directory {}", mDataDir);
      // If the UFS path is not valid, abort the test
      IOTaskResult result = new IOTaskResult();
      result.setParameters(mParameters);
      result.setBaseParameters(mBaseParameters);
      result.addError(ExceptionUtils.asPlainText(e));
      return result;
    }

    List<CompletableFuture<IOTaskResult>> futures = new ArrayList<>();
    final byte[] randomData = CommonUtils.randomBytes(BUFFER_SIZE);
    for (int i = 0; i < numThreads; i++) {
      final int idx = i;
      CompletableFuture<IOTaskResult> future = CompletableFuture.supplyAsync(() -> {
        IOTaskResult result = new IOTaskResult();
        result.setParameters(mParameters);
        result.setBaseParameters(mBaseParameters);
        long startTime = CommonUtils.getCurrentMs();

        String filePath = getFilePath(idx);
        LOG.debug("filePath={}, data to write={}", filePath, mParameters.mDataSize);

        long wroteBytes = 0;
        BufferedOutputStream outStream = null;
        try {
          outStream = new BufferedOutputStream(ufs.create(filePath));
          while (wroteBytes < ioSizeBytes) {
            long bytesToWrite = Math.min(ioSizeBytes - wroteBytes, BUFFER_SIZE);
            // bytesToWrite is bounded by BUFFER_SIZE, which is an integer
            outStream.write(randomData, 0, (int) bytesToWrite);
            wroteBytes += bytesToWrite;
          }
          outStream.flush();

          long endTime = CommonUtils.getCurrentMs();
          double duration = (endTime - startTime) / 1000.0; // convert to second
          IOTaskResult.Point p = new IOTaskResult.Point(IOTaskResult.IOMode.WRITE,
                  duration, wroteBytes);
          result.addPoint(p);
          LOG.debug("Write task finished {}", p);
        } catch (Exception e) {
          LOG.error("Failed to write to UFS: ", e);
          result.addError(e.getMessage());
        } finally {
          if (outStream != null) {
            try {
              outStream.close();
            } catch (IOException e) {
              LOG.warn("Failed to close stream to UFS: ", e);
              result.addError(e.getMessage());
            }
          }
        }

        LOG.debug("Thread {} file={}, IOBench result={}", Thread.currentThread().getName(),
                filePath, result);
        return result;
      }, pool);
      futures.add(future);
    }

    // Collect the result
    CompletableFuture[] cfs = futures.toArray(new CompletableFuture[0]);
    List<IOTaskResult> results = CompletableFuture.allOf(cfs)
            .thenApply(f -> futures.stream()
                    .map(CompletableFuture::join)
                    .collect(Collectors.toList())
            ).get();

    return IOTaskResult.reduceList(results);
  }
}
