// Copyright 2016 The Bazel 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 com.google.devtools.build.lib.remote;

import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.MetadataProvider;
import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
import com.google.devtools.build.lib.remote.DigestUtil.ActionKey;
import com.google.devtools.build.lib.remote.TreeNodeRepository.TreeNode;
import com.google.devtools.build.lib.remote.blobstore.SimpleBlobStore;
import com.google.devtools.build.lib.util.io.FileOutErr;
import com.google.devtools.build.lib.vfs.FileSystemUtils;
import com.google.devtools.build.lib.vfs.Path;
import com.google.devtools.remoteexecution.v1test.ActionResult;
import com.google.devtools.remoteexecution.v1test.Command;
import com.google.devtools.remoteexecution.v1test.Digest;
import com.google.devtools.remoteexecution.v1test.Directory;
import com.google.devtools.remoteexecution.v1test.DirectoryNode;
import com.google.devtools.remoteexecution.v1test.FileNode;
import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Collection;
import java.util.Map;

/**
 * A RemoteActionCache implementation that uses a concurrent map as a distributed storage for files
 * and action output.
 *
 * <p>The thread safety is guaranteed by the underlying map.
 *
 * <p>Note that this class is used from src/tools/remote.
 */
@ThreadSafe
public final class SimpleBlobStoreActionCache extends AbstractRemoteActionCache {
  private static final int MAX_BLOB_SIZE_FOR_INLINE = 10 * 1024;

  private final SimpleBlobStore blobStore;

  public SimpleBlobStoreActionCache(SimpleBlobStore blobStore, DigestUtil digestUtil) {
    super(digestUtil);
    this.blobStore = blobStore;
  }

  @Override
  public void ensureInputsPresent(
      TreeNodeRepository repository, Path execRoot, TreeNode root, Command command)
          throws IOException, InterruptedException {
    repository.computeMerkleDigests(root);
    uploadBlob(command.toByteArray());
    for (Directory directory : repository.treeToDirectories(root)) {
      uploadBlob(directory.toByteArray());
    }
    // TODO(ulfjack): Only upload files that aren't in the CAS yet?
    for (TreeNode leaf : repository.leaves(root)) {
      uploadFileContents(leaf.getActionInput(), execRoot, repository.getInputFileCache());
    }
  }

  public void downloadTree(Digest rootDigest, Path rootLocation)
      throws IOException, InterruptedException {
    FileSystemUtils.createDirectoryAndParents(rootLocation);
    Directory directory = Directory.parseFrom(downloadBlob(rootDigest));
    for (FileNode file : directory.getFilesList()) {
      downloadFile(
          rootLocation.getRelative(file.getName()), file.getDigest(), file.getIsExecutable(), null);
    }
    for (DirectoryNode child : directory.getDirectoriesList()) {
      downloadTree(child.getDigest(), rootLocation.getRelative(child.getName()));
    }
  }

  private Digest uploadFileContents(Path file) throws IOException, InterruptedException {
    Digest digest = digestUtil.compute(file);
    try (InputStream in = file.getInputStream()) {
      return uploadStream(digest, in);
    }
  }

  private Digest uploadFileContents(
      ActionInput input, Path execRoot, MetadataProvider inputCache)
          throws IOException, InterruptedException {
    if (input instanceof VirtualActionInput) {
      byte[] blob = ((VirtualActionInput) input).getBytes().toByteArray();
      return uploadBlob(blob, digestUtil.compute(blob));
    }
    try (InputStream in = execRoot.getRelative(input.getExecPathString()).getInputStream()) {
      return uploadStream(DigestUtil.getFromInputCache(input, inputCache), in);
    }
  }

  @Override
  public void upload(
      ActionKey actionKey,
      Path execRoot,
      Collection<Path> files,
      FileOutErr outErr,
      boolean uploadAction)
      throws IOException, InterruptedException {
    ActionResult.Builder result = ActionResult.newBuilder();
    upload(result, execRoot, files);
    if (outErr.getErrorPath().exists()) {
      Digest stderr = uploadFileContents(outErr.getErrorPath());
      result.setStderrDigest(stderr);
    }
    if (outErr.getOutputPath().exists()) {
      Digest stdout = uploadFileContents(outErr.getOutputPath());
      result.setStdoutDigest(stdout);
    }
    if (uploadAction) {
      blobStore.putActionResult(actionKey.getDigest().getHash(), result.build().toByteArray());
    }
  }

  public void upload(ActionResult.Builder result, Path execRoot, Collection<Path> files)
      throws IOException, InterruptedException {
    UploadManifest manifest = new UploadManifest(result, execRoot);
    manifest.addFiles(files);

    for (Map.Entry<Digest, Path> entry : manifest.getDigestToFile().entrySet()) {
      try (InputStream in = entry.getValue().getInputStream()) {
        uploadStream(entry.getKey(), in);
      }
    }

    for (Map.Entry<Digest, Chunker> entry : manifest.getDigestToChunkers().entrySet()) {
      uploadBlob(entry.getValue().next().getData().toByteArray(), entry.getKey());
    }
  }

  public void uploadOutErr(ActionResult.Builder result, byte[] stdout, byte[] stderr)
      throws IOException, InterruptedException {
    if (stdout.length <= MAX_BLOB_SIZE_FOR_INLINE) {
      result.setStdoutRaw(ByteString.copyFrom(stdout));
    } else if (stdout.length > 0) {
      result.setStdoutDigest(uploadBlob(stdout));
    }
    if (stderr.length <= MAX_BLOB_SIZE_FOR_INLINE) {
      result.setStderrRaw(ByteString.copyFrom(stderr));
    } else if (stderr.length > 0) {
      result.setStderrDigest(uploadBlob(stderr));
    }
  }

  public Digest uploadBlob(byte[] blob) throws IOException, InterruptedException {
    return uploadBlob(blob, digestUtil.compute(blob));
  }

  private Digest uploadBlob(byte[] blob, Digest digest) throws IOException, InterruptedException {
    return uploadStream(digest, new ByteArrayInputStream(blob));
  }

  public Digest uploadStream(Digest digest, InputStream in)
      throws IOException, InterruptedException {
    blobStore.put(digest.getHash(), digest.getSizeBytes(), in);
    return digest;
  }

  public boolean containsKey(Digest digest) throws IOException, InterruptedException {
    return blobStore.containsKey(digest.getHash());
  }

  @Override
  public ActionResult getCachedActionResult(ActionKey actionKey)
      throws IOException, InterruptedException {
    try {
      byte[] data = downloadActionResult(actionKey.getDigest());
      return ActionResult.parseFrom(data);
    } catch (InvalidProtocolBufferException | CacheNotFoundException e) {
      return null;
    }
  }

  private byte[] downloadActionResult(Digest digest) throws IOException, InterruptedException {
    if (digest.getSizeBytes() == 0) {
      return new byte[0];
    }
    // This unconditionally downloads the whole blob into memory!
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    boolean success = blobStore.getActionResult(digest.getHash(), out);
    if (!success) {
      throw new CacheNotFoundException(digest);
    }
    return out.toByteArray();
  }

  public void setCachedActionResult(ActionKey actionKey, ActionResult result)
      throws IOException, InterruptedException {
    blobStore.putActionResult(actionKey.getDigest().getHash(), result.toByteArray());
  }

  @Override
  public void close() {
    blobStore.close();
  }

  @Override
  protected void downloadBlob(Digest digest, Path dest) throws IOException, InterruptedException {
    try (OutputStream out = dest.getOutputStream()) {
      boolean success = blobStore.get(digest.getHash(), out);
      if (!success) {
        throw new CacheNotFoundException(digest);
      }
    }
  }

  @Override
  public byte[] downloadBlob(Digest digest) throws IOException, InterruptedException {
    if (digest.getSizeBytes() == 0) {
      return new byte[0];
    }
    try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
      boolean success = blobStore.get(digest.getHash(), out);
      if (!success) {
        throw new CacheNotFoundException(digest);
      }
      return out.toByteArray();
    }
  }
}