/*
 * Copyright (C) 2018 The Android Open Source Project
 *
 * 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.android.tools.build.bundletool.model.utils;

import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static java.util.function.Function.identity;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.MultimapBuilder;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Streams;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.stream.Collector;
import java.util.stream.Collectors;

/** Utility class providing custom {@link Collector}s. */
public final class CollectorUtils {

  /**
   * Returns a {@code Collector} accumulating entries into an {@code ImmutableListMultimap}.
   *
   * <p>The keys of the entries are the result of applying the provided mapping function while the
   * values are accumulated in the encounter order of the stream.
   */
  public static <T, K extends Comparable<K>>
      Collector<T, ?, ImmutableListMultimap<K, T>> groupingBySortedKeys(
          Function<? super T, ? extends K> keyFunction) {
    return groupingBySortedKeys(keyFunction, identity());
  }

  /**
   * Returns a {@code Collector} accumulating entries into an {@code ImmutableListMultimap}.
   *
   * <p>The keys of the entries are the result of applying the provided key mapping function while
   * the values are generated by applying the value mapping function and accumulated in the
   * encounter order of the stream.
   */
  public static <T, K extends Comparable<K>, V>
      Collector<T, ?, ImmutableListMultimap<K, V>> groupingBySortedKeys(
          Function<? super T, ? extends K> keyFunction,
          Function<? super T, ? extends V> valueFunction) {
    return Collectors.collectingAndThen(
        Multimaps.toMultimap(
            keyFunction, valueFunction, MultimapBuilder.treeKeys().arrayListValues()::<K, V>build),
        ImmutableListMultimap::copyOf);
  }

  public static <T, K> Collector<T, ?, ImmutableMap<K, ImmutableList<T>>> groupingByDeterministic(
      Function<? super T, ? extends K> keyFunction) {
    return Collectors.collectingAndThen(
        Collectors.groupingBy(keyFunction, LinkedHashMap::new, toImmutableList()),
        ImmutableMap::copyOf);
  }

  public static <T, K, D, A> Collector<T, ?, ImmutableMap<K, D>> groupingByDeterministic(
      Function<? super T, ? extends K> keyFunction, Collector<? super T, A, D> valueCollector) {
    return Collectors.collectingAndThen(
        Collectors.groupingBy(keyFunction, LinkedHashMap::new, valueCollector),
        ImmutableMap::copyOf);
  }

  /**
   * Returns a single immutable copy of the mappings in {@code map1} and {@code map2}. Conflicts are
   * resolved using the given {@code mergeFunction}.
   */
  public static <K, V> ImmutableMap<K, V> combineMaps(
      Map<K, V> map1, Map<K, V> map2, BinaryOperator<V> mergeFunction) {
    return Streams.concat(map1.entrySet().stream(), map2.entrySet().stream())
        .collect(toImmutableMap(Entry::getKey, Entry::getValue, mergeFunction));
  }

  private CollectorUtils() {}
}
