/*
 * Copyright 2018-present Facebook, 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.
 */

package com.facebook.buck.rules.modern.builders;

import com.facebook.buck.rules.modern.builders.Protocol.OutputDirectory;
import com.facebook.buck.util.RichStream;
import com.facebook.buck.util.function.ThrowingFunction;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.io.MoreFiles;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.UncheckedExecutionException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.function.Consumer;

/** Used for materialzing outputs from the CAS. */
public class OutputsMaterializer {
  private final AsyncBlobFetcher fetcher;
  private final Protocol protocol;

  public OutputsMaterializer(AsyncBlobFetcher fetcher, Protocol protocol) {
    this.fetcher = fetcher;
    this.protocol = protocol;
  }

  private static Path ensureParent(Path path) throws IOException {
    MoreFiles.createParentDirectories(path);
    return path;
  }

  /** Materialize the outputs of an action into a directory. */
  public void materialize(
      Collection<OutputDirectory> outputDirectories,
      Collection<Protocol.OutputFile> outputFiles,
      Path root)
      throws IOException {
    ImmutableList.Builder<ListenableFuture<Void>> pending = ImmutableList.builder();

    for (Protocol.OutputFile file : outputFiles) {
      Path path = root.resolve(file.getPath());
      ensureParent(path);
      ByteBuffer content = file.getContent();
      if (content != null) {
        try (FileChannel output =
            FileChannel.open(path, StandardOpenOption.WRITE, StandardOpenOption.CREATE)) {
          output.write(content);
        }
        setExecutable(file.getIsExecutable(), path);
      } else {
        pending.add(fetchAndMaterialize(file.getDigest(), file.getIsExecutable(), path));
      }
    }

    for (Protocol.OutputDirectory directory : outputDirectories) {
      Path dirRoot = root.resolve(directory.getPath());
      pending.add(
          Futures.transformAsync(
              fetcher.fetch(directory.getTreeDigest()),
              data -> {
                Protocol.Tree tree = protocol.parseTree(data);
                Map<Protocol.Digest, Protocol.Directory> childMap =
                    RichStream.from(tree.getChildrenList())
                        .collect(
                            ImmutableMap.toImmutableMap(
                                ThrowingFunction.asFunction(protocol::computeDigest),
                                child -> child));

                ImmutableList.Builder<ListenableFuture<Void>> pendingFilesBuilder =
                    ImmutableList.builder();
                materializeDirectory(childMap, tree.getRoot(), dirRoot, pendingFilesBuilder::add);
                return Futures.whenAllSucceed(pendingFilesBuilder.build()).call(() -> null);
              }));
    }

    try {
      Futures.whenAllSucceed(pending.build()).call(() -> null).get();
    } catch (InterruptedException | ExecutionException e) {
      throw new RuntimeException(e);
    }
  }

  private void materializeDirectory(
      Map<Protocol.Digest, Protocol.Directory> childMap,
      Protocol.Directory directory,
      Path root,
      Consumer<ListenableFuture<Void>> pendingWorkConsumer)
      throws IOException {
    Files.createDirectories(root);
    for (Protocol.DirectoryNode childNode : directory.getDirectoriesList()) {
      materializeDirectory(
          childMap,
          Preconditions.checkNotNull(childMap.get(childNode.getDigest())),
          root.resolve(childNode.getName()),
          pendingWorkConsumer);
    }

    for (Protocol.FileNode file : directory.getFilesList()) {
      pendingWorkConsumer.accept(
          fetchAndMaterialize(
              file.getDigest(), file.getIsExecutable(), root.resolve(file.getName())));
    }
  }

  private ListenableFuture<Void> fetchAndMaterialize(
      Protocol.Digest digest, boolean isExecutable, Path path) {
    return Futures.transform(
        fetcher.fetch(digest),
        data -> {
          try {
            try (FileOutputStream fileStream = new FileOutputStream(path.toFile())) {
              fileStream.getChannel().write(data);
            }
            if (isExecutable) {
              setExecutable(true, path);
            }
            return null;
          } catch (IOException e) {
            throw new UncheckedExecutionException(e);
          }
        });
  }

  private void setExecutable(boolean isExecutable, Path path) {
    if (isExecutable) {
      Preconditions.checkState(path.toFile().setExecutable(true));
    }
  }
}

@@ -769,65 +769,8 @@
 am;%0A
-import com.facebook.buck.util.function.ThrowingFunction;%0A
 impo
@@ -3359,44 +3359,106 @@
     
-ThrowingFunction.asFunction(
+d -%3E %7B%0A                                  try %7B%0A                                    return 
 protocol
 ::co
@@ -3453,18 +3453,17 @@
 protocol
-::
+.
 computeD
@@ -3463,25 +3463,225 @@
 omputeDigest
-)
+(d);%0A                                  %7D catch (IOException e) %7B%0A                                    throw new RuntimeException(e);%0A                                  %7D%0A                                %7D
 ,%0A          

Use a more efficient and cross compiler approach.