// Copyright 2019 The Nomulus Authors. All Rights Reserved.
//
// 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.

package google.registry.gradle.plugin;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Strings.isNullOrEmpty;
import static google.registry.gradle.plugin.GcsPluginUtils.readFilesWithEntryPoint;
import static google.registry.gradle.plugin.GcsPluginUtils.toByteArraySupplier;
import static google.registry.gradle.plugin.GcsPluginUtils.toNormalizedPath;
import static google.registry.gradle.plugin.GcsPluginUtils.uploadFileToGcs;
import static google.registry.gradle.plugin.GcsPluginUtils.uploadFilesToGcsMultithread;

import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.storage.Storage;
import com.google.cloud.storage.StorageOptions;
import com.google.common.collect.ImmutableMap;
import com.google.common.io.Files;
import google.registry.gradle.plugin.ProjectData.TaskData;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Optional;
import java.util.function.BiConsumer;
import java.util.function.Supplier;
import org.gradle.api.DefaultTask;
import org.gradle.api.Project;
import org.gradle.api.Task;
import org.gradle.api.reporting.DirectoryReport;
import org.gradle.api.reporting.Report;
import org.gradle.api.reporting.ReportContainer;
import org.gradle.api.reporting.Reporting;
import org.gradle.api.tasks.TaskAction;

/** A task that uploads the Reports generated by other tasks to GCS. */
public class ReportUploader extends DefaultTask {

  private static final SecureRandom SECURE_RANDOM = new SecureRandom();
  private static final ImmutableMap<String, BiConsumer<ReportUploader, String>> UPLOAD_FUNCTIONS =
      ImmutableMap.of(
          "file://", ReportUploader::saveResultsToLocalFolder,
          "gcs://", ReportUploader::uploadResultsToGcs);

  private final ArrayList<Task> tasks = new ArrayList<>();
  private final HashMap<String, StringBuilder> logs = new HashMap<>();
  private Project project;

  private String destination = null;
  private String credentialsFile = null;
  private String multithreadedUpload = null;

  /**
   * Sets the destination of the reports.
   *
   * <p>Currently supports two types of destinations:
   *
   * <ul>
   *   <li>file://[absulute local path], e.g. file:///tmp/buildOutputs/
   *   <li>gcs://[bucket name]/[optional path], e.g. gcs://my-bucket/buildOutputs/
   * </ul>
   */
  public void setDestination(String destination) {
    this.destination = destination;
  }

  public void setCredentialsFile(String credentialsFile) {
    this.credentialsFile = credentialsFile;
  }

  public void setMultithreadedUpload(String multithreadedUpload) {
    this.multithreadedUpload = multithreadedUpload;
  }

  /** Converts the given Gradle Project into a ProjectData. */
  private ProjectData createProjectData() {
    ProjectData.Builder builder =
        ProjectData.builder()
            .setName(project.getPath() + project.getName())
            .setDescription(
                Optional.ofNullable(project.getDescription()).orElse("[No description available]"))
            .setGradleVersion(project.getGradle().getGradleVersion())
            .setProjectProperties(project.getGradle().getStartParameter().getProjectProperties())
            .setSystemProperties(project.getGradle().getStartParameter().getSystemPropertiesArgs())
            .setTasksRequested(project.getGradle().getStartParameter().getTaskNames());

    Path rootDir = toNormalizedPath(project.getRootDir());
    tasks.stream()
        .filter(task -> task.getState().getExecuted() || task.getState().getUpToDate())
        .map(task -> createTaskData(task, rootDir))
        .forEach(builder.tasksBuilder()::add);
    return builder.build();
  }

  /**
   * Converts a Gradle Task into a TaskData.
   *
   * @param rootDir the root directory of the main Project - used to get the relative path of any
   *     Task files.
   */
  private TaskData createTaskData(Task task, Path rootDir) {
    TaskData.State state =
        task.getState().getFailure() != null
            ? TaskData.State.FAILURE
            : task.getState().getUpToDate() ? TaskData.State.UP_TO_DATE : TaskData.State.SUCCESS;
    String log = logs.get(task.getPath()).toString();

    TaskData.Builder builder =
        TaskData.builder()
            .setState(state)
            .setUniqueName(task.getPath())
            .setDescription(
                Optional.ofNullable(task.getDescription()).orElse("[No description available]"));
    if (!log.isEmpty()) {
      builder.setLog(toByteArraySupplier(log));
    }

    Reporting<? extends ReportContainer<? extends Report>> reporting = asReporting(task);

    if (reporting != null) {
      // This Task is also a Reporting task! It has a destination file/directory for every supported
      // format.
      // Add the files for each of the formats into the ReportData.
      reporting
          .getReports()
          .getAsMap()
          .forEach(
              (type, report) -> {
                File destination = report.getDestination();
                // The destination could be a file, or a directory. If it's a directory - the Report
                // could have created multiple files - and we need to know to which one of those to
                // link.
                //
                // If we're lucky, whoever implemented the Report made sure to extend
                // DirectoryReport, which gives us the entry point to all the files.
                //
                // This isn't guaranteed though, as it depends on the implementer.
                Optional<File> entryPointHint =
                    destination.isDirectory() && (report instanceof DirectoryReport)
                        ? Optional.ofNullable(((DirectoryReport) report).getEntryPoint())
                        : Optional.empty();
                builder
                    .reportsBuilder()
                    .put(type, readFilesWithEntryPoint(destination, entryPointHint, rootDir));
              });
    }
    return builder.build();
  }

  private FilesWithEntryPoint generateFilesToUpload() {
    ProjectData projectData = createProjectData();
    CoverPageGenerator coverPageGenerator = new CoverPageGenerator(projectData);
    return coverPageGenerator.getFilesToUpload();
  }

  @TaskAction
  void uploadResults() {
    System.out.format("ReportUploader: destination= '%s'\n", destination);

    try {

      if (isNullOrEmpty(destination)) {
        System.out.format("ReportUploader: no destination given, skipping...\n");
        return;
      }

      for (String key : UPLOAD_FUNCTIONS.keySet()) {
        if (destination.startsWith(key)) {
          UPLOAD_FUNCTIONS.get(key).accept(this, destination.substring(key.length()));
          return;
        }
      }
      System.out.format(
          "ReportUploader: given destination '%s' doesn't start with one of %s."
              + " Defaulting to saving in /tmp\n",
          destination, UPLOAD_FUNCTIONS.keySet());
      saveResultsToLocalFolder("/tmp/");
    } catch (Throwable e) {
      System.out.format("ReportUploader: Encountered error %s\n", e);
      e.printStackTrace(System.out);
      System.out.format("ReportUploader: skipping upload\n");
    }
  }

  private void saveResultsToLocalFolder(String absoluteFolderName) {
    Path folder = Paths.get(absoluteFolderName, createUniqueFolderName());
    checkArgument(
        folder.isAbsolute(),
        "Local files destination must be an absolute path, but is %s",
        absoluteFolderName);
    FilesWithEntryPoint filesToUpload = generateFilesToUpload();
    System.out.format(
        "ReportUploader: going to save %s files to %s\n", filesToUpload.files().size(), folder);
    filesToUpload
        .files()
        .forEach((path, dataSupplier) -> saveFile(folder.resolve(path), dataSupplier));
    System.out.format(
        "ReportUploader: report saved to file://%s\n", folder.resolve(filesToUpload.entryPoint()));
  }

  private void saveFile(Path path, Supplier<byte[]> dataSupplier) {
    File dir = path.getParent().toFile();
    if (!dir.isDirectory()) {
      checkState(dir.mkdirs(), "Couldn't create directory %s", dir);
    }
    try {
      Files.write(dataSupplier.get(), path.toFile());
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    }
  }

  private void uploadResultsToGcs(String destination) {
    checkArgument(
        !destination.isEmpty(), "destination must include at least the bucket name, but is empty");
    Path bucketWithFolder = Paths.get(destination, createUniqueFolderName());
    String bucket = bucketWithFolder.getName(0).toString();
    Path folder = bucketWithFolder.subpath(1, bucketWithFolder.getNameCount());

    StorageOptions.Builder storageOptions = StorageOptions.newBuilder();
    if (!isNullOrEmpty(credentialsFile)) {
      try {
        storageOptions.setCredentials(
            GoogleCredentials.fromStream(new FileInputStream(credentialsFile)));
      } catch (IOException e) {
        throw new UncheckedIOException(e);
      }
    }
    Storage storage = storageOptions.build().getService();

    FilesWithEntryPoint filesToUpload = generateFilesToUpload();

    System.out.format(
        "ReportUploader: going to upload %s files to %s/%s\n",
        filesToUpload.files().size(), bucket, folder);
    if ("yes".equals(multithreadedUpload)) {
      System.out.format("ReportUploader: multi-threaded upload\n");
      uploadFilesToGcsMultithread(storage, bucket, folder, filesToUpload.files());
    } else {
      System.out.format("ReportUploader: single threaded upload\n");
      filesToUpload
          .files()
          .forEach(
              (path, dataSupplier) -> {
                System.out.format("ReportUploader: Uploading %s\n", path);
                uploadFileToGcs(storage, bucket, folder.resolve(path), dataSupplier);
              });
    }
    System.out.format(
        "ReportUploader: report uploaded to https://storage.googleapis.com/%s/%s\n",
        bucket, folder.resolve(filesToUpload.entryPoint()));
  }

  void setProject(Project project) {
    this.project = project;

    for (Project subProject : project.getAllprojects()) {
      subProject.getTasks().all(this::addTask);
    }
  }

  private void addTask(Task task) {
    if (task instanceof ReportUploader) {
      return;
    }
    tasks.add(task);
    StringBuilder log = new StringBuilder();
    checkArgument(
        !logs.containsKey(task.getPath()),
        "Multiple tasks with the same .getPath()=%s",
        task.getPath());
    logs.put(task.getPath(), log);
    task.getLogging().addStandardOutputListener(output -> log.append(output));
    task.getLogging().addStandardErrorListener(output -> log.append(output));
    task.finalizedBy(this);
  }

  @SuppressWarnings("unchecked")
  private static Reporting<? extends ReportContainer<? extends Report>> asReporting(Task task) {
    if (task instanceof Reporting) {
      return (Reporting<? extends ReportContainer<? extends Report>>) task;
    }
    return null;
  }

  private String createUniqueFolderName() {
    return String.format(
        "%h-%h-%h-%h",
        SECURE_RANDOM.nextInt(),
        SECURE_RANDOM.nextInt(),
        SECURE_RANDOM.nextInt(),
        SECURE_RANDOM.nextInt());
  }
}
