

package com.facebook.buck.parser;

import com.facebook.buck.log.Logger;
import com.facebook.buck.model.BuildTarget;
import com.facebook.buck.model.BuildTargetException;
import com.facebook.buck.model.UnflavoredBuildTarget;
import com.facebook.buck.parser.thrift.BuildFileEnvProperty;
import com.facebook.buck.parser.thrift.RemoteDaemonicCellState;
import com.facebook.buck.rules.Cell;
import com.facebook.buck.util.MoreCollectors;
import com.facebook.buck.util.ObjectMappers;
import com.facebook.buck.util.concurrent.AutoCloseableLock;
import com.facebook.buck.util.concurrent.AutoCloseableReadWriteUpdateLock;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Preconditions;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.MapDifference;
import com.google.common.collect.Maps;
import com.google.common.collect.SetMultimap;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.concurrent.GuardedBy;

class DaemonicCellState {

  private static final Logger LOG = Logger.get(DaemonicCellState.class);

  private class CacheImpl<T> implements PipelineNodeCache.Cache<BuildTarget, T> {

    @GuardedBy("rawAndComputedNodesLock")
    public final ConcurrentMapCache<BuildTarget, T> allComputedNodes =
        new ConcurrentMapCache<>(parsingThreads);

    @Override
    public Optional<T> lookupComputedNode(Cell cell, BuildTarget target)
        throws BuildTargetException {
      try (AutoCloseableLock readLock = rawAndComputedNodesLock.readLock()) {
        return Optional.ofNullable(allComputedNodes.getIfPresent(target));
      }
    }

    @Override
    public T putComputedNodeIfNotPresent(Cell cell, BuildTarget target, T targetNode)
        throws BuildTargetException {
      try (AutoCloseableLock writeLock = rawAndComputedNodesLock.writeLock()) {
        T updatedNode = allComputedNodes.putIfAbsentAndGet(target, targetNode);
        Preconditions.checkState(
            allRawNodeTargets.contains(target.getUnflavoredBuildTarget()),
            "Added %s to computed nodes, which isn't present in raw nodes",
            target);
        if (updatedNode.equals(targetNode)) {
          targetsCornucopia.put(target.getUnflavoredBuildTarget(), target);
        }
        return updatedNode;
      }
    }
  }

  private final Path cellRoot;
  private final Optional<String> cellCanonicalName;
  private AtomicReference<Cell> cell;

  @GuardedBy("rawAndComputedNodesLock")
  private final SetMultimap<Path, Path> buildFileDependents;

  @GuardedBy("rawAndComputedNodesLock")
  private final SetMultimap<UnflavoredBuildTarget, BuildTarget> targetsCornucopia;

  @GuardedBy("rawAndComputedNodesLock")
  private final Map<Path, ImmutableMap<String, Optional<String>>> buildFileEnv;

  @GuardedBy("rawAndComputedNodesLock")
  private final ConcurrentMapCache<Path, ImmutableSet<Map<String, Object>>> allRawNodes;
  // Tracks all targets in `allRawNodes`.  Used to verify that every target in `allComputedNodes`
  // is also in `allRawNodes`, as we use the latter for bookkeeping invalidations.
  @GuardedBy("rawAndComputedNodesLock")
  private final Set<UnflavoredBuildTarget> allRawNodeTargets;

  @GuardedBy("rawAndComputedNodesLock")
  private final ConcurrentMap<Class<?>, CacheImpl<?>> typedNodeCaches;

  private final AutoCloseableReadWriteUpdateLock rawAndComputedNodesLock;
  private final int parsingThreads;

  DaemonicCellState(Cell cell, int parsingThreads) {
    this.cell = new AtomicReference<>(cell);
    this.parsingThreads = parsingThreads;
    this.cellRoot = cell.getRoot();
    this.cellCanonicalName = cell.getCanonicalName();
    this.buildFileDependents = HashMultimap.create();
    this.targetsCornucopia = HashMultimap.create();
    this.buildFileEnv = new HashMap<>();
    this.allRawNodes = new ConcurrentMapCache<>(parsingThreads);
    this.allRawNodeTargets = new HashSet<>();
    this.typedNodeCaches = Maps.newConcurrentMap();
    this.rawAndComputedNodesLock = new AutoCloseableReadWriteUpdateLock();
  }

  // TODO(mzlee): Only needed for invalidateBasedOn which does not have access to cell metadata
  Cell getCell() {
    return Preconditions.checkNotNull(cell.get());
  }

  Path getCellRoot() {
    return cellRoot;
  }

  @SuppressWarnings("unchecked")
  public <T> CacheImpl<T> getOrCreateCache(Class<T> type) {
    try (AutoCloseableLock updateLock = rawAndComputedNodesLock.updateLock()) {
      CacheImpl<?> cache = typedNodeCaches.get(type);
      if (cache == null) {
        try (AutoCloseableLock writeLock = rawAndComputedNodesLock.writeLock()) {
          cache = new CacheImpl<>();
          typedNodeCaches.put(type, cache);
        }
      }
      return (CacheImpl<T>) cache;
    }
  }

  @SuppressWarnings("unchecked")
  public <T> CacheImpl<T> getCache(Class<T> type) {
    try (AutoCloseableLock readLock = rawAndComputedNodesLock.readLock()) {
      return (CacheImpl<T>) typedNodeCaches.get(type);
    }
  }

  Optional<ImmutableSet<Map<String, Object>>> lookupRawNodes(Path buildFile) {
    try (AutoCloseableLock readLock = rawAndComputedNodesLock.readLock()) {
      return Optional.ofNullable(allRawNodes.getIfPresent(buildFile));
    }
  }

  ImmutableSet<Map<String, Object>> putRawNodesIfNotPresentAndStripMetaEntries(
      final Path buildFile,
      final ImmutableSet<Map<String, Object>> withoutMetaIncludes,
      final ImmutableSet<Path> dependentsOfEveryNode,
      ImmutableMap<String, Optional<String>> env) {
    try (AutoCloseableLock writeLock = rawAndComputedNodesLock.writeLock()) {
      ImmutableSet<Map<String, Object>> updated =
          allRawNodes.putIfAbsentAndGet(buildFile, withoutMetaIncludes);
      for (Map<String, Object> node : updated) {
        allRawNodeTargets.add(
            RawNodeParsePipeline.parseBuildTargetFromRawRule(
                cellRoot, cellCanonicalName, node, buildFile));
      }
      buildFileEnv.put(buildFile, env);
      if (updated == withoutMetaIncludes) {
        // We now know all the nodes. They all implicitly depend on everything in
        // the "dependentsOfEveryNode" set.
        for (Path dependent : dependentsOfEveryNode) {
          buildFileDependents.put(dependent, buildFile);
        }
      }
      return updated;
    }
  }

  int invalidatePath(Path path) {
    try (AutoCloseableLock writeLock = rawAndComputedNodesLock.writeLock()) {
      int invalidatedRawNodes = 0;
      ImmutableSet<Map<String, Object>> rawNodes = allRawNodes.getIfPresent(path);
      if (rawNodes != null) {
        // Increment the counter
        invalidatedRawNodes = rawNodes.size();
        for (Map<String, Object> rawNode : rawNodes) {
          UnflavoredBuildTarget target =
              RawNodeParsePipeline.parseBuildTargetFromRawRule(
                  cellRoot, cellCanonicalName, rawNode, path);
          LOG.debug("Invalidating target for path %s: %s", path, target);
          for (CacheImpl<?> cache : typedNodeCaches.values()) {
            cache.allComputedNodes.invalidateAll(targetsCornucopia.get(target));
          }
          targetsCornucopia.removeAll(target);
          allRawNodeTargets.remove(target);
        }
        allRawNodes.invalidate(path);
      }

      // We may have been given a file that other build files depend on. Iteratively remove those.
      Iterable<Path> dependents = buildFileDependents.get(path);
      LOG.debug("Invalidating dependents for path %s: %s", path, dependents);
      for (Path dependent : dependents) {
        if (dependent.equals(path)) {
          continue;
        }
        invalidatedRawNodes += invalidatePath(dependent);
      }
      buildFileDependents.removeAll(path);
      buildFileEnv.remove(path);

      return invalidatedRawNodes;
    }
  }

  Optional<MapDifference<String, String>> invalidateIfEnvHasChanged(Cell cell, Path buildFile) {
    // Invalidate if env vars have changed.
    ImmutableMap<String, Optional<String>> usedEnv;
    try (AutoCloseableLock readLock = rawAndComputedNodesLock.readLock()) {
      usedEnv = buildFileEnv.get(buildFile);
    }
    if (usedEnv == null) {
      this.cell.set(cell);
      return Optional.empty();
    }
    for (Map.Entry<String, Optional<String>> ent : usedEnv.entrySet()) {
      Optional<String> value =
          Optional.ofNullable(cell.getBuckConfig().getEnvironment().get(ent.getKey()));
      if (!value.equals(ent.getValue())) {
        LOG.verbose("invalidating for env change: %s (%s != %s)", buildFile, value, ent.getValue());
        invalidatePath(buildFile);
        this.cell.set(cell);
        return Optional.of(
            Maps.difference(
                value.map(v -> ImmutableMap.of(ent.getKey(), v)).orElse(ImmutableMap.of()),
                ent.getValue()
                    .map(v -> ImmutableMap.of(ent.getKey(), v))
                    .orElse(ImmutableMap.of())));
      }
    }
    return Optional.empty();
  }

  private Map<String, String> getAllRawNodesForSerialisation() throws IOException {
    Map<String, String> result = new HashMap<>();
    Path root = getCellRoot();
    ObjectMapper objectMapper = new ObjectMapper();
    for (Path path : allRawNodes.keySet()) {
      ImmutableSet<Map<String, Object>> v = allRawNodes.getIfPresent(path);
      if (v != null) {
        result.put(root.relativize(path).toString(), objectMapper.writeValueAsString(v));
      }
    }
    return result;
  }

  private Map<String, List<String>> getBuildFileDependentsForSerialisation() {
    Map<String, List<String>> result = new HashMap<>();
    Path root = getCellRoot();
    Map<Path, Collection<Path>> view = buildFileDependents.asMap();
    for (Path path : view.keySet()) {
      Collection<Path> paths = view.get(path);
      if (paths != null) {
        ImmutableList<String> pathList =
            paths
                .stream()
                .map(v -> root.relativize(v).toString())
                .collect(MoreCollectors.toImmutableList());
        result.put(root.relativize(path).toString(), pathList);
      }
    }

    return result;
  }

  private Map<String, Map<String, BuildFileEnvProperty>> getBuildFileEnvForSerialisation() {
    Map<String, Map<String, BuildFileEnvProperty>> result = new HashMap<>();
    for (Path path : buildFileEnv.keySet()) {
      ImmutableMap.Builder<String, BuildFileEnvProperty> buildFileEnvValuesMapBuilder =
          ImmutableMap.builder();
      ImmutableMap<String, Optional<String>> values = buildFileEnv.get(path);
      if (values != null) {
        values.forEach(
            (k, v) -> {
              BuildFileEnvProperty prop = new BuildFileEnvProperty();
              v.ifPresent(prop::setValue);
              buildFileEnvValuesMapBuilder.put(k, prop);
            });
      }
      result.put(path.toString(), buildFileEnvValuesMapBuilder.build());
    }

    return result;
  }

  RemoteDaemonicCellState serialise() throws IOException {
    RemoteDaemonicCellState result = new RemoteDaemonicCellState();
    try (AutoCloseableLock readLock = rawAndComputedNodesLock.readLock()) {
      result.allRawNodesJsons = getAllRawNodesForSerialisation();
      result.buildFileDependents = getBuildFileDependentsForSerialisation();
      result.buildFileEnv = getBuildFileEnvForSerialisation();
    }
    return result;
  }

  static DaemonicCellState deserialise(
      RemoteDaemonicCellState remote, Cell cell, int parsingThreads) throws IOException {
    DaemonicCellState daemonicCellState = new DaemonicCellState(cell, parsingThreads);
    Path root = cell.getRoot();
    for (String pathString : remote.buildFileDependents.keySet()) {
      Path key = root.resolve(pathString);
      remote
          .buildFileDependents
          .get(pathString)
          .forEach(
              valuePathString -> {
                daemonicCellState.buildFileDependents.put(key, root.resolve(valuePathString));
              });
    }

    for (String pathString : remote.buildFileEnv.keySet()) {
      Path key = Paths.get(pathString);
      Map<String, BuildFileEnvProperty> remoteValues = remote.buildFileEnv.get(pathString);
      ImmutableMap.Builder<String, Optional<String>> builder = ImmutableMap.builder();
      for (String k : remoteValues.keySet()) {
        BuildFileEnvProperty prop = remoteValues.get(k);
        if (prop != null && !prop.isSetValue()) {
          builder.put(k, Optional.empty());
        } else {
          builder.put(k, Optional.of(prop.value));
        }
      }
      daemonicCellState.buildFileEnv.put(key, builder.build());
    }

    for (String pathString : remote.allRawNodesJsons.keySet()) {
      String json = remote.allRawNodesJsons.get(pathString);
      ImmutableSet<Map<String, Object>> deserialisedRawNodes =
          ObjectMappers.readValue(json, new TypeReference<ImmutableSet<Map<String, Object>>>() {});
      Path key = root.resolve(pathString);
      daemonicCellState.allRawNodes.putIfAbsentAndGet(key, deserialisedRawNodes);
      deserialisedRawNodes.forEach(
          rawNode -> {
            daemonicCellState.allRawNodeTargets.add(
                UnflavoredBuildTarget.of(
                    root,
                    cell.getCanonicalName(),
                    "//" + rawNode.get("buck.base_path"),
                    (String) rawNode.get("name")));
          });
    }

    return daemonicCellState;
  }
}
