// Copyright 2021 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.bazel.bzlmod;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.google.devtools.build.lib.analysis.BlazeDirectories;
import com.google.devtools.build.lib.cmdline.Label;
import com.google.devtools.build.lib.cmdline.PackageIdentifier;
import com.google.devtools.build.lib.cmdline.RepositoryMapping;
import com.google.devtools.build.lib.cmdline.RepositoryName;
import com.google.devtools.build.lib.cmdline.StarlarkThreadContext;
import com.google.devtools.build.lib.events.ExtendedEventHandler;
import com.google.devtools.build.lib.packages.NoSuchPackageException;
import com.google.devtools.build.lib.packages.Rule;
import com.google.devtools.build.lib.packages.RuleClass;
import com.google.devtools.build.lib.packages.RuleFactory.InvalidRuleException;
import com.google.devtools.build.lib.packages.StarlarkNativeModule.ExistingRulesShouldBeNoOp;
import java.util.LinkedHashMap;
import java.util.Map;
import javax.annotation.Nullable;
import net.starlark.java.eval.Dict;
import net.starlark.java.eval.EvalException;
import net.starlark.java.eval.NoneType;
import net.starlark.java.eval.Starlark;
import net.starlark.java.eval.StarlarkInt;
import net.starlark.java.eval.StarlarkList;
import net.starlark.java.eval.StarlarkSemantics;
import net.starlark.java.eval.StarlarkThread;
import net.starlark.java.syntax.Location;

/**
 * A context object that should be stored in a {@link StarlarkThread} for use during module
 * extension evaluation.
 */
public final class ModuleExtensionEvalStarlarkThreadContext extends StarlarkThreadContext {
  @Override
  public void storeInThread(StarlarkThread thread) {
    super.storeInThread(thread);
    // The following is just a hack; see documentation there for an explanation.
    thread.setThreadLocal(ExistingRulesShouldBeNoOp.class, new ExistingRulesShouldBeNoOp());
  }

  @Nullable
  public static ModuleExtensionEvalStarlarkThreadContext fromOrNull(StarlarkThread thread) {
    StarlarkThreadContext ctx = thread.getThreadLocal(StarlarkThreadContext.class);
    return ctx instanceof ModuleExtensionEvalStarlarkThreadContext c ? c : null;
  }

  record RepoRuleCall(
      RuleClass ruleClass,
      Dict<String, Object> kwargs,
      Location location,
      ImmutableList<StarlarkThread.CallStackEntry> callStack) {}

  private final ModuleExtensionId extensionId;
  private final String repoPrefix;
  private final PackageIdentifier basePackageId;
  private final RepositoryMapping baseRepoMapping;
  private final ImmutableMap<String, RepositoryName> repoOverrides;
  private final BlazeDirectories directories;
  private final ExtendedEventHandler eventHandler;
  private final Map<String, RepoRuleCall> deferredRepos = new LinkedHashMap<>();

  public ModuleExtensionEvalStarlarkThreadContext(
      ModuleExtensionId extensionId,
      String repoPrefix,
      PackageIdentifier basePackageId,
      RepositoryMapping baseRepoMapping,
      ImmutableMap<String, RepositoryName> repoOverrides,
      RepositoryMapping mainRepoMapping,
      BlazeDirectories directories,
      ExtendedEventHandler eventHandler) {
    super(() -> mainRepoMapping);
    this.extensionId = extensionId;
    this.repoPrefix = repoPrefix;
    this.basePackageId = basePackageId;
    this.baseRepoMapping = baseRepoMapping;
    this.repoOverrides = repoOverrides;
    this.directories = directories;
    this.eventHandler = eventHandler;
  }

  /**
   * Records a call to a repo rule that should be created at the end of the module extension
   * evaluation.
   */
  @SuppressWarnings("unchecked")
  public void lazilyCreateRepo(
      StarlarkThread thread, Dict<String, Object> kwargs, RuleClass ruleClass)
      throws EvalException {
    Object nameValue = kwargs.getOrDefault("name", Starlark.NONE);
    if (!(nameValue instanceof String name)) {
      throw Starlark.errorf(
          "expected string for attribute 'name', got '%s'", Starlark.type(nameValue));
    }
    RepositoryName.validateUserProvidedRepoName(name);
    RepoRuleCall conflict = deferredRepos.get(name);
    if (conflict != null) {
      throw Starlark.errorf(
          "A repo named %s is already generated by this module extension at %s",
          name, conflict.location());
    }
    deferredRepos.put(
        name,
        new RepoRuleCall(
            ruleClass,
            // The extension may mutate the values of the kwargs after this function returns.
            (Dict<String, Object>) deepCloneAttrValue(kwargs),
            thread.getCallerLocation(),
            thread.getCallStack()));
  }

  /**
   * Evaluates the repo rule calls recorded by {@link #lazilyCreateRepo} and returns all repos
   * generated by the extension. The key is the "internal name" (as specified by the extension) of
   * the repo, and the value is the {@link RepoSpec}.
   */
  public ImmutableMap<String, RepoSpec> createRepos(StarlarkSemantics starlarkSemantics)
      throws EvalException, InterruptedException {
    // LINT.IfChange
    // Make it possible to refer to extension repos in the label attributes of another extension
    // repo. Wrapping a label in Label(...) ensures that it is evaluated with respect to the
    // containing module's repo mapping instead.
    ImmutableMap.Builder<String, RepositoryName> entries = ImmutableMap.builder();
    entries.putAll(baseRepoMapping.entries());
    entries.putAll(
        Maps.asMap(
            deferredRepos.keySet(),
            apparentName -> RepositoryName.createUnvalidated(repoPrefix + apparentName)));
    entries.putAll(repoOverrides);
    RepositoryMapping fullRepoMapping =
        RepositoryMapping.create(entries.buildKeepingLast(), baseRepoMapping.ownerRepo());
    // LINT.ThenChange(//src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionRepoMappingEntriesFunction.java)

    ImmutableMap.Builder<String, RepoSpec> repoSpecs = ImmutableMap.builder();
    for (var entry : deferredRepos.entrySet()) {
      String name = entry.getKey();
      RepoRuleCall repoRuleCall = entry.getValue();
      try {
        String prefixedName = repoPrefix + name;
        Rule rule =
            BzlmodRepoRuleCreator.createRule(
                basePackageId,
                fullRepoMapping,
                directories,
                starlarkSemantics,
                eventHandler,
                repoRuleCall.callStack,
                repoRuleCall.ruleClass,
                Maps.transformEntries(
                    repoRuleCall.kwargs, (k, v) -> k.equals("name") ? prefixedName : v));

        Map<String, Object> attributes =
            Maps.filterKeys(
                Maps.transformEntries(repoRuleCall.kwargs, (k, v) -> rule.getAttr(k)),
                k -> !k.equals("name"));
        var attributesValue = AttributeValues.create(attributes);
        AttributeValues.validateAttrs(
            attributesValue,
            String.format("in the extension '%s'", extensionId),
            String.format("%s '%s'", rule.getRuleClass(), name));
        repoSpecs.put(
            name,
            new RepoSpec(
                new RepoRuleId(
                    repoRuleCall.ruleClass.getRuleDefinitionEnvironmentLabel(),
                    repoRuleCall.ruleClass.getName()),
                attributesValue));
      } catch (EvalException e) {
        throw e.withCallStack(repoRuleCall.callStack);
      } catch (InvalidRuleException | NoSuchPackageException e) {
        throw new EvalException(e).withCallStack(repoRuleCall.callStack);
      }
    }
    return repoSpecs.buildOrThrow();
  }

  /**
   * Deep-clones a potentially mutable Starlark object that is a valid repo rule attribute.
   * Immutable (sub-)objects are not cloned.
   */
  private static Object deepCloneAttrValue(Object x) throws EvalException {
    if (x instanceof NoneType
        || x instanceof Boolean
        || x instanceof StarlarkInt
        || x instanceof String
        || x instanceof Label) {
      return x;
    }
    // Mutable Starlark values have to be cloned deeply.
    if (x instanceof Dict<?, ?> dict) {
      Dict.Builder<Object, Object> newDict = Dict.builder();
      for (Map.Entry<?, ?> e : dict.entrySet()) {
        newDict.put(deepCloneAttrValue(e.getKey()), deepCloneAttrValue(e.getValue()));
      }
      return newDict.buildImmutable();
    }
    if (x instanceof Iterable<?> iterable) {
      ImmutableList.Builder<Object> newList = ImmutableList.builder();
      for (Object item : iterable) {
        newList.add(deepCloneAttrValue(item));
      }
      return StarlarkList.immutableCopyOf(newList.build());
    }
    throw Starlark.errorf(
        "unexpected Starlark value: %s (of type %s)", Starlark.repr(x), Starlark.type(x));
  }
}
