// Copyright 2023 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.rules.starlarkdocextract;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Verify.verify;
import static com.google.common.base.Verify.verifyNotNull;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.devtools.build.lib.packages.ImplicitOutputsFunction.fromTemplates;
import static com.google.devtools.build.lib.packages.Type.BOOLEAN;
import static com.google.devtools.build.lib.packages.Type.STRING_LIST;
import static java.util.stream.Collectors.partitioningBy;

import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.io.ByteSource;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.actions.MutableActionGraph.ActionConflictException;
import com.google.devtools.build.lib.analysis.ConfiguredTarget;
import com.google.devtools.build.lib.analysis.FileProvider;
import com.google.devtools.build.lib.analysis.RuleConfiguredTargetBuilder;
import com.google.devtools.build.lib.analysis.RuleConfiguredTargetFactory;
import com.google.devtools.build.lib.analysis.RuleContext;
import com.google.devtools.build.lib.analysis.Runfiles;
import com.google.devtools.build.lib.analysis.RunfilesProvider;
import com.google.devtools.build.lib.analysis.actions.BinaryFileWriteAction;
import com.google.devtools.build.lib.analysis.actions.FileWriteAction;
import com.google.devtools.build.lib.cmdline.BazelModuleContext;
import com.google.devtools.build.lib.cmdline.Label;
import com.google.devtools.build.lib.cmdline.RepositoryMapping;
import com.google.devtools.build.lib.cmdline.RepositoryName;
import com.google.devtools.build.lib.collect.nestedset.NestedSet;
import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
import com.google.devtools.build.lib.collect.nestedset.Order;
import com.google.devtools.build.lib.packages.ImplicitOutputsFunction.SafeImplicitOutputsFunction;
import com.google.devtools.build.lib.profiler.Profiler;
import com.google.devtools.build.lib.profiler.SilentCloseable;
import com.google.devtools.build.lib.skyframe.BzlLoadFunction.BzlLoadFailedException;
import com.google.devtools.build.lib.skyframe.BzlLoadValue;
import com.google.devtools.build.lib.skyframe.RepositoryMappingValue;
import com.google.devtools.build.lib.skyframe.RepositoryMappingValue.RepositoryMappingResolutionException;
import com.google.devtools.build.skydoc.rendering.LabelRenderer;
import com.google.devtools.build.skydoc.rendering.proto.StardocOutputProtos.ModuleInfo;
import com.google.devtools.build.skyframe.SkyFunction;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.protobuf.TextFormat;
import java.io.IOException;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Predicate;
import javax.annotation.Nullable;
import net.starlark.java.eval.Module;

/** Implementation of the {@code starlark_doc_extract} rule. */
public class StarlarkDocExtract implements RuleConfiguredTargetFactory {
  static final String SRC_ATTR = "src";
  static final String DEPS_ATTR = "deps";
  static final String SYMBOL_NAMES_ATTR = "symbol_names";
  static final String RENDER_MAIN_REPO_NAME = "render_main_repo_name";
  static final SafeImplicitOutputsFunction BINARYPROTO_OUT = fromTemplates("%{name}.binaryproto");
  static final SafeImplicitOutputsFunction TEXTPROTO_OUT = fromTemplates("%{name}.textproto");

  @Override
  @Nullable
  public ConfiguredTarget create(RuleContext ruleContext)
      throws ActionConflictException, InterruptedException, RuleErrorException {
    RepositoryMappingValue mainRepositoryMappingValue = getMainRepositoryMappingValue(ruleContext);
    RepositoryMapping repositoryMapping = mainRepositoryMappingValue.getRepositoryMapping();
    Module module = loadModule(ruleContext, repositoryMapping);
    if (module == null) {
      // Skyframe restart
      verify(
          ruleContext.getAnalysisEnvironment().getSkyframeEnv().valuesMissing()
              && !ruleContext.hasErrors());
      return null;
    }
    verifyModuleDeps(ruleContext, module, repositoryMapping);
    Optional<String> mainRepoName = Optional.empty();
    if (ruleContext.attributes().get(RENDER_MAIN_REPO_NAME, BOOLEAN)) {
      mainRepoName = mainRepositoryMappingValue.getAssociatedModuleName();
      if (mainRepoName.isEmpty()) {
        mainRepoName = Optional.of(ruleContext.getWorkspaceName());
      }
    }
    ModuleInfo moduleInfo =
        getModuleInfo(ruleContext, module, new LabelRenderer(repositoryMapping, mainRepoName));

    NestedSet<Artifact> filesToBuild =
        new NestedSetBuilder<Artifact>(Order.STABLE_ORDER)
            .add(createBinaryProtoOutput(ruleContext, moduleInfo))
            .build();
    // Textproto output isn't in filesToBuild: we want to create it only if explicitly requested.
    createTextProtoOutput(ruleContext, moduleInfo);

    return new RuleConfiguredTargetBuilder(ruleContext)
        .setFilesToBuild(filesToBuild)
        .addProvider(
            RunfilesProvider.class,
            RunfilesProvider.simple(
                new Runfiles.Builder(ruleContext.getWorkspaceName())
                    .addRunfiles(ruleContext, RunfilesProvider.DEFAULT_RUNFILES)
                    .addTransitiveArtifacts(filesToBuild)
                    .build()))
        .build();
  }

  /**
   * Loads the Starlark module from the source file given by the rule's {@code src} attribute.
   *
   * @throws RuleErrorException and reports an error in the rule if the {@code src} attribute refers
   *     to multiple or zero files, a generated file, or a source file which cannot be loaded or
   *     parsed
   * @return the module object, or null on Skyframe restart
   */
  @Nullable
  private static Module loadModule(RuleContext ruleContext, RepositoryMapping repositoryMapping)
      throws RuleErrorException, InterruptedException {
    try (SilentCloseable c = Profiler.instance().profile("BzlDocDump.loadModule")) {
      // Note attr schema validates that src is a .bzl or .scl file.
      Label label = getSourceFileLabel(ruleContext, SRC_ATTR, repositoryMapping);

      // Note getSkyframeEnv() cannot be null while creating a configured target.
      SkyFunction.Environment env = ruleContext.getAnalysisEnvironment().getSkyframeEnv();

      BzlLoadValue bzlLoadValue;
      try {
        // TODO(b/276733504): support loading modules in @_builtins
        bzlLoadValue =
            (BzlLoadValue)
                env.getValueOrThrow(BzlLoadValue.keyForBuild(label), BzlLoadFailedException.class);
      } catch (BzlLoadFailedException e) {
        ruleContext.attributeError(SRC_ATTR, e.getMessage());
        throw new RuleErrorException(e);
      }
      if (bzlLoadValue == null) {
        // Skyframe restart
        return null;
      }
      return bzlLoadValue.getModule();
    }
  }

  /**
   * Retrieves the label of the singular source artifact from a given attribute. Note that we can't
   * simply use {@code ruleContext.attributes().get(attrName, LABEL)} because that does not resolve
   * aliases and filegroups.
   *
   * @throws RuleErrorException if the source is not a singular source artifact, meaning its label
   *     cannot be used as a label for a Starlark load()
   */
  private static Label getSourceFileLabel(
      RuleContext ruleContext, String attrName, RepositoryMapping repositoryMapping)
      throws RuleErrorException {
    Artifact artifact = ruleContext.getPrerequisiteArtifact(attrName);
    ruleContext.assertNoErrors();
    // If ruleContext.getPrerequisiteArtifact() set no errors, we know artifact != null
    if (!artifact.isSourceArtifact()) {
      RuleErrorException error =
          new RuleErrorException(
              String.format(
                  "%s is not a source file and cannot be loaded in Starlark",
                  formatDerivedArtifact(artifact, repositoryMapping)));
      ruleContext.attributeError(attrName, error.getMessage());
      throw error;
    }
    return verifyNotNull(artifact.getOwner());
  }

  private static String formatDerivedArtifact(
      Artifact artifact, RepositoryMapping repositoryMapping) {
    checkArgument(!artifact.isSourceArtifact());
    return String.format(
        "%s (generated by rule %s)",
        artifact.getRepositoryRelativePath(),
        artifact.getOwner().getDisplayForm(repositoryMapping));
  }

  /**
   * Verifies that the module's transitive loads are a subset of the source artifacts in
   * files-to-build of the rule's deps.
   *
   * @throws RuleErrorException if that is not the case.
   */
  // TODO(https://github.com/bazelbuild/bazel/issues/18599): to avoid flattening deps, we could use
  // either (a) a new, native bzl_library-like rule that verifies strict deps, or (b) a new native
  // aspect that verifies strict deps for the existing bzl_library rule. Ideally, however, we ought
  // to get rid of the deps attribute (and the need to verify it) altogether; that requires new
  // dependency machinery for `bazel query` to use the Starlark load graph for collecting the
  // dependencies of starlark_doc_extract's src.
  private static void verifyModuleDeps(
      RuleContext ruleContext, Module module, RepositoryMapping repositoryMapping)
      throws RuleErrorException {
    // Note attr schema validates that deps are .bzl or .scl files.
    Map<Boolean, ImmutableSet<Artifact>> flattenedDepsPartitionedByIsSource =
        ruleContext.getPrerequisites(DEPS_ATTR).stream()
            // TODO(https://github.com/bazelbuild/bazel/issues/18599): we are using FileProvider
            // instead of StarlarkLibraryInfo only because StarlarkLibraryInfo is defined in
            // bazel_skylib, not natively in Bazel.
            .flatMap(dep -> dep.getProvider(FileProvider.class).getFilesToBuild().toList().stream())
            .collect(partitioningBy(Artifact::isSourceArtifact, toImmutableSet()));
    // bzl_library targets may contain both source artifacts and derived artifacts (e.g. generated
    // .bzl files for tests); only the source artifacts can be load()-ed by Bazel.
    ImmutableSet<Artifact> flattenedDepsSourceArtifacts =
        flattenedDepsPartitionedByIsSource.getOrDefault(true, ImmutableSet.of());
    ImmutableSet<Artifact> flattenedDepsDerivedArtifacts =
        flattenedDepsPartitionedByIsSource.getOrDefault(false, ImmutableSet.of());

    ImmutableList<String> topmostUnknownLoads =
        getTopmostUnknownLoads(
            module,
            flattenedDepsSourceArtifacts.stream()
                .map(artifact -> verifyNotNull(artifact.getOwner()))
                .collect(toImmutableSet()),
            repositoryMapping);

    if (!topmostUnknownLoads.isEmpty()) {
      StringBuilder errorMessageBuilder =
          new StringBuilder("missing bzl_library targets for Starlark module(s) ")
              .append(Joiner.on(", ").join(topmostUnknownLoads));
      if (!flattenedDepsDerivedArtifacts.isEmpty()) {
        // TODO(arostovtsev): we ought to print only the derived artifacts having the same
        // root-relative path as topmostUnknownLoads.
        errorMessageBuilder
            .append("\nNote the following are generated file(s) and cannot be loaded in Starlark: ")
            .append(
                Joiner.on(", ")
                    .join(
                        flattenedDepsDerivedArtifacts.stream()
                            .map(artifact -> formatDerivedArtifact(artifact, repositoryMapping))
                            .iterator()));
      }
      RuleErrorException error = new RuleErrorException(errorMessageBuilder.toString());
      ruleContext.attributeError(DEPS_ATTR, error.getMessage());
      throw error;
    }
  }

  /**
   * Finds the topmost modules that are transitively loaded by the given module but not mentioned in
   * the given set of known modules, and returns these modules' display forms.
   *
   * <p>Unknown modules that are only referenced by other unknown modules are not included.
   */
  private static ImmutableList<String> getTopmostUnknownLoads(
      Module module, ImmutableSet<Label> knownModules, RepositoryMapping repositoryMapping) {
    ImmutableList.Builder<String> unknown = ImmutableList.builder();
    Set<Label> visited = new LinkedHashSet<>();
    BazelModuleContext.visitLoadGraphRecursively(
        BazelModuleContext.of(module).loads(),
        label -> {
          if (!visited.add(label)) {
            return false;
          }
          if (!knownModules.contains(label)) {
            unknown.add(label.getDisplayForm(repositoryMapping));
            return false;
          }
          return true;
        });
    return unknown.build();
  }

  /** Returns the main repository's repo mapping value. */
  private static RepositoryMappingValue getMainRepositoryMappingValue(RuleContext ruleContext)
      throws RuleErrorException, InterruptedException {
    RepositoryMappingValue repositoryMappingValue;
    try {
      repositoryMappingValue =
          (RepositoryMappingValue)
              ruleContext
                  .getAnalysisEnvironment()
                  .getSkyframeEnv()
                  .getValueOrThrow(
                      RepositoryMappingValue.key(RepositoryName.MAIN),
                      RepositoryMappingResolutionException.class);
    } catch (RepositoryMappingResolutionException e) {
      ruleContext.ruleError(e.getMessage());
      throw new RuleErrorException(e);
    }
    verifyNotNull(repositoryMappingValue);
    return repositoryMappingValue;
  }

  private static ModuleInfo getModuleInfo(
      RuleContext ruleContext, Module module, LabelRenderer labelRenderer)
      throws RuleErrorException {
    ModuleInfo moduleInfo;
    try {
      moduleInfo =
          new ModuleInfoExtractor(getWantedSymbolPredicate(ruleContext), labelRenderer)
              .extractFrom(module);
    } catch (ModuleInfoExtractor.ExtractionException e) {
      ruleContext.ruleError(e.getMessage());
      throw new RuleErrorException(e);
    }
    return moduleInfo;
  }

  private static Predicate<String> getWantedSymbolPredicate(RuleContext ruleContext) {
    ImmutableList<String> symbolNames =
        ImmutableList.copyOf(ruleContext.attributes().get(SYMBOL_NAMES_ATTR, STRING_LIST));
    if (symbolNames.isEmpty()) {
      return name -> true;
    } else {
      return symbolNames::contains;
    }
  }

  private static Artifact createBinaryProtoOutput(RuleContext ruleContext, ModuleInfo moduleInfo)
      throws InterruptedException {
    Artifact binaryProtoOutput = ruleContext.getImplicitOutputArtifact(BINARYPROTO_OUT);
    ruleContext.registerAction(
        new BinaryFileWriteAction(
            ruleContext.getActionOwner(),
            binaryProtoOutput,
            ByteSource.wrap(moduleInfo.toByteArray()),
            /* makeExecutable= */ false));
    return binaryProtoOutput;
  }

  @Nullable
  @CanIgnoreReturnValue
  private static Artifact createTextProtoOutput(RuleContext ruleContext, ModuleInfo moduleInfo)
      throws InterruptedException, RuleErrorException {
    Artifact textProtoOutput = ruleContext.getImplicitOutputArtifact(TEXTPROTO_OUT);

    StringBuilder textprotoBuilder = new StringBuilder();
    try {
      TextFormat.printer().print(moduleInfo, textprotoBuilder);
    } catch (IOException e) {
      ruleContext.ruleError(e.getMessage());
      throw new RuleErrorException(e);
    }
    ruleContext.registerAction(
        FileWriteAction.create(
            ruleContext,
            textProtoOutput,
            textprotoBuilder.toString(),
            /* makeExecutable= */ false));
    return textProtoOutput;
  }
}