/**
 * Copyright (c) 2019. Qubole Inc
 * Licensed 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. See accompanying LICENSE file.
 */
package com.qubole.rubix.client.robotframework;

import com.qubole.rubix.core.MockCachingFileSystem;
import com.qubole.rubix.spi.BookKeeperFactory;
import com.qubole.rubix.spi.CacheUtil;
import com.qubole.rubix.spi.RetryingPooledBookkeeperClient;
import com.qubole.rubix.spi.thrift.BlockLocation;
import com.qubole.rubix.spi.thrift.CacheStatusRequest;
import org.apache.commons.io.FileUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.Path;
import org.apache.thrift.shaded.TException;
import org.apache.thrift.shaded.transport.TTransportException;

import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

import static com.qubole.rubix.spi.utils.DataSizeUnits.BYTES;

public class BookKeeperClientRFLibrary
{
  private final BookKeeperFactory factory = new BookKeeperFactory();
  private final Configuration conf = new Configuration();
  private static final String FILE_SCHEME = "file:";

  /**
   * Read data from a given file into the BookKeeper cache using the BookKeeper Thrift API.
   *
   * @param readRequest The read request to execute.
   * @return True if the data was read into the cache correctly, false otherwise.
   */
  public boolean cacheDataUsingBookKeeperServerCall(TestClientReadRequest readRequest) throws IOException, TException
  {
    try (RetryingPooledBookkeeperClient client = createBookKeeperClient()) {
      return client.readData(
          getPathWithFileScheme(readRequest.getRemotePath()),
          readRequest.getReadStart(),
          readRequest.getReadLength(),
          readRequest.getFileLength(),
          readRequest.getLastModified(),
          readRequest.getClusterType());
    }
  }

  /**
   * Read data into the BookKeeper cache from multiple concurrent clients using the BookKeeper Thrift API.
   *
   * @param numThreads    The number of threads to use for concurrent execution.
   * @param readRequests  The read requests to concurrently execute.
   * @return True if all read requests succeeded, false otherwise.
   */
  public boolean concurrentlyCacheDataUsingBookKeeperServerCall(int numThreads,
                                                                boolean staggerRequests,
                                                                List<TestClientReadRequest> readRequests) throws TException, InterruptedException, ExecutionException
  {
    final List<Callable<Boolean>> tasks = new ArrayList<>();
    for (final TestClientReadRequest request : readRequests) {
      final Callable<Boolean> callable = new Callable<Boolean>()
      {
        @Override
        public Boolean call() throws Exception
        {
          return cacheDataUsingBookKeeperServerCall(request);
        }
      };
      tasks.add(callable);
    }

    final List<Future<Boolean>> results = executeConcurrentTasks(numThreads, tasks, staggerRequests);
    final boolean didAllSucceed = didConcurrentDataDownloadSucceed(results);
    return didAllSucceed;
  }

  /**
   * Read data from a given file into the BookKeeper cache using a client caching file system.
   *
   * @param readRequest The read request to execute.
   * @return True if the data was read into the cache correctly, false otherwise.
   */
  public boolean cacheDataUsingClientFileSystem(TestClientReadRequest readRequest) throws IOException, TException, URISyntaxException
  {
    try (FSDataInputStream inputStream = createFSInputStream(getPathWithFileScheme(readRequest.getRemotePath()), readRequest.getReadLength())) {
      final int readSize = inputStream.read(
          new byte[readRequest.getReadLength()],
          (int) readRequest.getReadStart(),
          readRequest.getReadLength());
      return readSize == readRequest.getReadLength();
    }
  }

  /**
   * Read data into the BookKeeper cache from multiple concurrent clients using a client caching file system.
   *
   * @param numThreads    The number of threads to use for concurrent execution.
   * @param readRequests  The read requests to concurrently execute.
   * @return True if all read requests succeeded, false otherwise.
   */
  public boolean concurrentlyCacheDataUsingClientFileSystem(int numThreads,
                                                            boolean staggerRequests,
                                                            List<TestClientReadRequest> readRequests) throws TException, InterruptedException, ExecutionException
  {
    final List<Callable<Boolean>> tasks = new ArrayList<>();
    for (final TestClientReadRequest request : readRequests) {
      final Callable<Boolean> callable = new Callable<Boolean>()
      {
        @Override
        public Boolean call() throws Exception
        {
          return cacheDataUsingClientFileSystem(request);
        }
      };
      tasks.add(callable);
    }

    final List<Future<Boolean>> results = executeConcurrentTasks(numThreads, tasks, staggerRequests);
    final boolean didAllSucceed = didConcurrentDataDownloadSucceed(results);
    return didAllSucceed;
  }

  /**
   * Get the cache status for blocks in a particular file.
   *
   * @param request The request specifying the blocks & file for which to check the status.
   * @return A list of {@link BlockLocation}s detailing the locations of the specified blocks.
   */
  public List<BlockLocation> getCacheStatus(TestClientStatusRequest request) throws IOException, TException
  {
    try (RetryingPooledBookkeeperClient client = createBookKeeperClient()) {
      return client.getCacheStatus(new CacheStatusRequest(
          getPathWithFileScheme(request.getRemotePath()),
          request.getFileLength(),
          request.getLastModified(),
          request.getStartBlock(),
          request.getEndBlock()).setClusterType(request.getClusterType()));
    }
  }

  /**
   * Get the current cache metrics from the BookKeeper server.
   *
   * @return A map of metrics describing cache statistics and interactions.
   */
  public Map<String, Double> getCacheMetrics() throws IOException, TException
  {
    try (RetryingPooledBookkeeperClient client = createBookKeeperClient()) {
      return client.getCacheMetrics();
    }
  }

  /**
   * Get the combined size of all configured cache directories.
   *
   * @param dirPath     The root path for the cache directories
   * @param dirSuffix   The cache directory suffix.
   * @param numDisks    The expected number of cache disks.
   * @return The size of the cache in MB.
   */
  public long getCacheDirSizeMb(String dirPath, String dirSuffix, int numDisks)
  {
    long cacheSize = 0;
    for (int disk = 0; disk < numDisks; disk++) {
      long cacheDirSize = FileUtils.sizeOfDirectory(new File(dirPath + disk + dirSuffix));
      cacheSize += cacheDirSize;
    }

    return BYTES.toMB(cacheSize);
  }

  /**
   * Generate a file to be used during testing.
   *
   * @param filename  The name of the file.
   * @param size      The size of the file in bytes.
   * @throws IOException if an error occurs while writing to the specified file.
   */
  public void generateTestFile(String filename, long size) throws IOException
  {
    String content = generateContent(size);
    Files.write(Paths.get(filename), content.getBytes());
  }

  /**
   * Generate a metadata file to be used for testing situations where metadata exists without its matching cache file.
   *
   * @param filename  The name of the file.
   * @throws IOException if an error occurs while writing to the specified file.
   */
  public String generateTestMDFile(String filename) throws IOException
  {
    String mdPath = CacheUtil.getMetadataFilePath(filename, conf);
    // Certain tests require a non-empty metadata file.
    Files.write(Paths.get(mdPath), "0101010101".getBytes());
    return mdPath;
  }

  /**
   * Create a read request to be executed by the BookKeeper server.
   *
   * @param remotePath    The remote path location.
   * @param readStart     The block to start reading from.
   * @param readLength    The amount of data to read.
   * @param fileLength    The length of the file.
   * @param lastModified  The time at which the file was last modified.
   * @param clusterType   The type id of cluster being used.
   * @return The read request.
   */
  public TestClientReadRequest createTestClientReadRequest(String remotePath,
                                                           long readStart,
                                                           int readLength,
                                                           long fileLength,
                                                           long lastModified,
                                                           int clusterType)
  {
    return new TestClientReadRequest(remotePath, readStart, readLength, fileLength, lastModified, clusterType);
  }

  /**
   * Create a status request to be executed by the BookKeeper server.
   *
   * @param remotePath    The remote path location.
   * @param fileLength    The length of the file.
   * @param lastModified  The time at which the file was last modified.
   * @param startBlock    The start of the block range to check.
   * @param endBlock      The end of the block range to check.
   * @param clusterType   The type id of cluster being used.
   * @return The status request.
   */
  public TestClientStatusRequest createTestClientStatusRequest(String remotePath,
                                                               long fileLength,
                                                               long lastModified,
                                                               long startBlock,
                                                               long endBlock,
                                                               int clusterType)
  {
    return new TestClientStatusRequest(remotePath, fileLength, lastModified, startBlock, endBlock, clusterType);
  }

  /**
   * Generate a random character string.
   *
   * @param size  The size of the string to generate.
   * @return A random string of the specified size.
   */
  private String generateContent(long size)
  {
    StringBuilder builder = new StringBuilder();
    Random random = new Random();

    for (int i = 0; i < size; i++) {
      char randomChar = (char) (33 + random.nextInt(93));
      builder.append(randomChar);
    }
    return builder.toString();
  }

  /**
   * Initializes the Hadoop configuration for the library.
   *
   * @param configurationOptions The options to set.
   */
  public void initializeLibraryConfiguration(Map<String, String> configurationOptions)
  {
    conf.clear();
    for (final Map.Entry<String, String> option : configurationOptions.entrySet()) {
      conf.set(option.getKey(), option.getValue());
    }
  }

  /**
   * Create a client for interacting to a BookKeeper server.
   *
   * @return The BookKeeper client.
   * @throws TTransportException if an error occurs when trying to connect to the BookKeeper server.
   */
  private RetryingPooledBookkeeperClient createBookKeeperClient() throws TTransportException
  {
    return factory.createBookKeeperClient(conf);
  }

  /**
   * Creates & initializes an input stream for executing client caching.
   *
   * @param remotePath  The path of the file to cache.
   * @return the input stream for the file.
   * @throws IOException if an error occurs when initializing the file system.
   */
  private FSDataInputStream createFSInputStream(String remotePath, int readLength) throws IOException
  {
    final MockCachingFileSystem mockFS = new MockCachingFileSystem();
    mockFS.initialize(URI.create(remotePath), conf);
    return mockFS.open(new Path(remotePath), readLength);
  }

  /**
   * Execute multiple tasks concurrently.
   *
   * @param numThreads   The number of available threads for concurrent execution.
   * @param tasks        The tasks to execute.
   * @param staggerTasks If true, add delay between task submissions.
   * @param <T>          The return type of the task.
   * @return A list of results for each task executed.
   * @throws InterruptedException if task execution is interrupted.
   */
  private <T> List<Future<T>> executeConcurrentTasks(int numThreads, List<Callable<T>> tasks, boolean staggerTasks) throws InterruptedException
  {
    final ExecutorService service = Executors.newFixedThreadPool(numThreads);
    List<Future<T>> futures = new ArrayList<>();

    if (staggerTasks) {
      // Necessary to preserve order of requests for certain tests.
      for (Callable<T> task : tasks) {
        futures.add(service.submit(task));
        Thread.sleep(100);
      }
    }
    else {
      futures = service.invokeAll(tasks);
    }
    return futures;
  }

  /**
   * Verify the results of read requests executed concurrently.
   *
   * @param readRequestResults   The results to verify.
   * @return True if all read requests succeeded, false otherwise.
   */
  private boolean didConcurrentDataDownloadSucceed(List<Future<Boolean>> readRequestResults) throws ExecutionException, InterruptedException
  {
    boolean didAllSucceed = true;
    for (final Future<Boolean> result : readRequestResults) {
      final Boolean didRead = result.get();
      didAllSucceed &= didRead;
    }
    return didAllSucceed;
  }

  /**
   * Watch the cache directory for changes and verify state of cached files.
   *
   * @param cacheDir     The directory to watch.
   * @param maxWaitTime  The maximum amount of time to wait for file events.
   * @param requests     The read requests describing the files to be watched for.
   * @return True if all expected files have been
   * @throws ExecutionException
   * @throws InterruptedException
   */
  public boolean watchCache(final String cacheDir, final int maxWaitTime, final List<TestClientReadRequest> requests) throws ExecutionException, InterruptedException
  {
    Future<Boolean> watcherResult = startCacheWatcher(cacheDir, requests, maxWaitTime);
    boolean didCache = watcherResult.get();
    return didCache;
  }

  /**
   * Watch the cache directory for changes and verify state of cached files.
   *
   * @param cacheDir     The directory to watch.
   * @param maxWaitTime  The maximum amount of time to wait for file events.
   * @param requests     The read requests describing the files to be watched for.
   * @return The {@link Future} for the result of the cache watcher.
   */
  private Future<Boolean> startCacheWatcher(final String cacheDir, final List<TestClientReadRequest> requests, final int maxWaitTime)
  {
    return Executors.newSingleThreadExecutor().submit(new Callable<Boolean>()
    {
      @Override
      public Boolean call() throws Exception
      {
        CacheWatcher watcher = new CacheWatcher(conf, Paths.get(cacheDir), maxWaitTime);
        return watcher.watchForCacheFiles(requests);
      }
    });
  }

  /**
   * Add the file scheme to the provided path for proper execution with the BookKeeper server.
   *
   * @param path  The path to update.
   * @return The provided path with the file scheme.
   */
  private String getPathWithFileScheme(String path)
  {
    return FILE_SCHEME + path;
  }
}