/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.beam.sdk.values;

import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.beam.sdk.transforms.Contextful;
import org.apache.beam.sdk.transforms.ProcessFunction;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.checkerframework.checker.nullness.qual.NonNull;

/**
 * A utility class for creating {@link TypeDescriptor} objects for different types, such as Java
 * primitive types, containers and {@link KV KVs} of other {@link TypeDescriptor} objects, and
 * extracting type variables of parameterized types (e.g. extracting the {@code OutputT} type
 * variable of a {@code DoFn<InputT, OutputT>}).
 */
public class TypeDescriptors {
  /**
   * The {@link TypeDescriptor} for Boolean. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Boolean&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Boolean
   */
  public static TypeDescriptor<Boolean> booleans() {
    return new TypeDescriptor<Boolean>() {};
  }

  /**
   * The {@link TypeDescriptor} for Double. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Double&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Double
   */
  public static TypeDescriptor<Double> doubles() {
    return new TypeDescriptor<Double>() {};
  }

  /**
   * The {@link TypeDescriptor} for Float. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Float&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Float
   */
  public static TypeDescriptor<Float> floats() {
    return new TypeDescriptor<Float>() {};
  }

  /**
   * The {@link TypeDescriptor} for Integer. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Integer&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Integer
   */
  public static TypeDescriptor<Integer> integers() {
    return new TypeDescriptor<Integer>() {};
  }

  /**
   * The {@link TypeDescriptor} for Long. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Long&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Long
   */
  public static TypeDescriptor<Long> longs() {
    return new TypeDescriptor<Long>() {};
  }

  /**
   * The {@link TypeDescriptor} for Short. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Short&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Short
   */
  public static TypeDescriptor<Short> shorts() {
    return new TypeDescriptor<Short>() {};
  }

  /**
   * The {@link TypeDescriptor} for BigDecimal. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;BigDecimal&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for BigDecimal
   */
  public static TypeDescriptor<BigDecimal> bigdecimals() {
    return new TypeDescriptor<BigDecimal>() {};
  }

  /**
   * The {@link TypeDescriptor} for BigInteger. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;BigInteger&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for BigInteger
   */
  public static TypeDescriptor<BigInteger> bigintegers() {
    return new TypeDescriptor<BigInteger>() {};
  }

  /**
   * The {@link TypeDescriptor} for {@link Row}. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Row&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Row
   */
  public static TypeDescriptor<Row> rows() {
    return new TypeDescriptor<Row>() {};
  }

  /**
   * The {@link TypeDescriptor} for String. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;String&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for String
   */
  public static TypeDescriptor<String> strings() {
    return new TypeDescriptor<String>() {};
  }

  /**
   * The {@link TypeDescriptor} for Character. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Character&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Character
   */
  public static TypeDescriptor<Character> characters() {
    return new TypeDescriptor<Character>() {};
  }

  /**
   * The {@link TypeDescriptor} for Byte. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Byte&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for Byte
   */
  public static TypeDescriptor<Byte> bytes() {
    return new TypeDescriptor<Byte>() {};
  }

  /**
   * The {@link TypeDescriptor} for nulls/Void. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Void&gt;() {};
   * </pre>
   *
   * @return A {@link TypeDescriptor} for nulls/Void
   */
  public static TypeDescriptor<Void> nulls() {
    return new TypeDescriptor<Void>() {};
  }

  /**
   * The {@link TypeDescriptor} for {@link KV}. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;KV&lt;K,V&gt;&gt;() {};
   * </pre>
   *
   * <p>Example of use:
   *
   * <pre>{@code
   * PCollection<String> words = ...;
   * PCollection<KV<String, String>> words = words.apply(FlatMapElements
   *         .into(TypeDescriptors.kv(TypeDescriptors.strings(), TypeDescriptors.strings()))
   *         .via(...));
   * }</pre>
   *
   * @param key The {@link TypeDescriptor} for the key
   * @param value The {@link TypeDescriptor} for the value
   * @return A {@link TypeDescriptor} for {@link KV}
   */
  public static <K, V> TypeDescriptor<KV<K, V>> kvs(
      TypeDescriptor<K> key, TypeDescriptor<V> value) {
    return new TypeDescriptor<KV<K, V>>() {}.where(new TypeParameter<K>() {}, key)
        .where(new TypeParameter<V>() {}, value);
  }

  /**
   * The {@link TypeDescriptor} for {@link Set}. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Set&lt;E&gt;&gt;() {};
   * </pre>
   *
   * <p>Example of use:
   *
   * <pre>{@code
   * PCollection<String> words = ...;
   * PCollection<Set<String>> words = words.apply(FlatMapElements
   *         .into(TypeDescriptors.sets(TypeDescriptors.strings()))
   *         .via(...));
   * }</pre>
   *
   * @param element The {@link TypeDescriptor} for the set
   * @return A {@link TypeDescriptor} for {@link Set}
   */
  public static <T> TypeDescriptor<Set<T>> sets(TypeDescriptor<T> element) {
    return new TypeDescriptor<Set<T>>() {}.where(new TypeParameter<T>() {}, element);
  }

  /** The {@link TypeDescriptor} for {@link Map}. */
  public static <K, V> TypeDescriptor<Map<K, V>> maps(
      TypeDescriptor<K> keyType, TypeDescriptor<V> valueType) {
    return new TypeDescriptor<Map<K, V>>() {}.where(new TypeParameter<K>() {}, keyType)
        .where(new TypeParameter<V>() {}, valueType);
  }

  /**
   * The {@link TypeDescriptor} for {@link List}. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;List&lt;E&gt;&gt;() {};
   * </pre>
   *
   * <p>Example of use:
   *
   * <pre>{@code
   * PCollection<String> words = ...;
   * PCollection<List<String>> words = words.apply(FlatMapElements
   *         .into(TypeDescriptors.lists(TypeDescriptors.strings()))
   *         .via(...));
   * }</pre>
   *
   * @param element The {@link TypeDescriptor} for the list
   * @return A {@link TypeDescriptor} for {@link List}
   */
  public static <T> TypeDescriptor<List<T>> lists(TypeDescriptor<T> element) {
    return new TypeDescriptor<List<T>>() {}.where(new TypeParameter<T>() {}, element);
  }

  /**
   * The {@link TypeDescriptor} for {@link Iterable}. This is the equivalent of:
   *
   * <pre>
   * new TypeDescriptor&lt;Iterable&lt;E&gt;&gt;() {};
   * </pre>
   *
   * <p>Example of use:
   *
   * <pre>{@code
   * PCollection<String> words = ...;
   * PCollection<Iterable<String>> words = words.apply(FlatMapElements
   *         .into(TypeDescriptors.iterables(TypeDescriptors.strings()))
   *         .via(...));
   * }</pre>
   *
   * @param iterable The {@link TypeDescriptor} for the iterable
   * @return A {@link TypeDescriptor} for {@link Iterable}
   */
  public static <T> TypeDescriptor<Iterable<T>> iterables(TypeDescriptor<T> iterable) {
    return new TypeDescriptor<Iterable<T>>() {}.where(new TypeParameter<T>() {}, iterable);
  }

  public static TypeDescriptor<Void> voids() {
    return new TypeDescriptor<Void>() {};
  }

  /**
   * A helper interface for use with {@link #extractFromTypeParameters(Object, Class,
   * TypeVariableExtractor)}.
   */
  public interface TypeVariableExtractor<InputT, OutputT> {}

  /**
   * Extracts a type from the actual type parameters of a parameterized class, subject to Java type
   * erasure. The type to extract is specified in a way that is safe w.r.t. changing the type
   * signature of the parameterized class, as opposed to specifying the name or index of a type
   * variable.
   *
   * <p>Example of use:
   *
   * <pre>{@code
   * class Foo<BarT> {
   *   private ProcessFunction<BarT, String> fn;
   *
   *   TypeDescriptor<BarT> inferBarTypeDescriptorFromFn() {
   *     return TypeDescriptors.extractFromTypeParameters(
   *       fn,
   *       ProcessFunction.class,
   *       // The actual type of "fn" is matched against the input type of the extractor,
   *       // and the obtained values of type variables of the superclass are substituted
   *       // into the output type of the extractor.
   *       new TypeVariableExtractor<ProcessFunction<BarT, String>, BarT>() {});
   *   }
   * }
   * }</pre>
   *
   * @param instance The object being analyzed
   * @param supertype Parameterized superclass of interest
   * @param extractor A class for specifying the type to extract from the supertype
   * @return A {@link TypeDescriptor} for the actual value of the result type of the extractor,
   *     potentially containing unresolved type variables if the type was erased.
   */
  @SuppressWarnings("unchecked")
  public static <T, V> TypeDescriptor<V> extractFromTypeParameters(
      @NonNull T instance, Class<? super T> supertype, TypeVariableExtractor<T, V> extractor) {
    return extractFromTypeParameters(
        (TypeDescriptor<T>) TypeDescriptor.of(instance.getClass()), supertype, extractor);
  }

  /**
   * Like {@link #extractFromTypeParameters(Object, Class, TypeVariableExtractor)}, but takes a
   * {@link TypeDescriptor} of the instance being analyzed rather than the instance itself.
   */
  @SuppressWarnings("unchecked")
  public static <T, V> TypeDescriptor<V> extractFromTypeParameters(
      TypeDescriptor<T> type, Class<? super T> supertype, TypeVariableExtractor<T, V> extractor) {
    // Get the type signature of the extractor, e.g.
    // TypeVariableExtractor<ProcessFunction<BarT, String>, BarT>
    TypeDescriptor<TypeVariableExtractor<T, V>> extractorSupertype =
        (TypeDescriptor<TypeVariableExtractor<T, V>>)
            TypeDescriptor.of(extractor.getClass()).getSupertype(TypeVariableExtractor.class);

    // Get the actual type argument, e.g. ProcessFunction<BarT, String>
    Type inputT = ((ParameterizedType) extractorSupertype.getType()).getActualTypeArguments()[0];

    // Get the actual supertype of the type being analyzed, hopefully with all type parameters
    // resolved, e.g. ProcessFunction<Integer, String>
    TypeDescriptor<? super T> supertypeDescriptor = type.getSupertype(supertype);

    // Substitute actual supertype into the extractor, e.g.
    // TypeVariableExtractor<ProcessFunction<Integer, String>, Integer>
    TypeDescriptor<TypeVariableExtractor<T, V>> extractorT =
        extractorSupertype.where(inputT, supertypeDescriptor.getType());

    // Get output of the extractor.
    Type outputT = ((ParameterizedType) extractorT.getType()).getActualTypeArguments()[1];
    return (TypeDescriptor<V>) TypeDescriptor.of(outputT);
  }

  /**
   * Returns a type descriptor for the input of the given {@link ProcessFunction}, subject to Java
   * type erasure: may contain unresolved type variables if the type was erased.
   */
  public static <InputT, OutputT> TypeDescriptor<InputT> inputOf(
      ProcessFunction<InputT, OutputT> fn) {
    return extractFromTypeParameters(
        fn,
        ProcessFunction.class,
        new TypeVariableExtractor<ProcessFunction<InputT, OutputT>, InputT>() {});
  }

  /** Binary compatibility adapter for {@link #inputOf(ProcessFunction)}. */
  public static <InputT, OutputT> TypeDescriptor<InputT> inputOf(
      SerializableFunction<InputT, OutputT> fn) {
    return inputOf((ProcessFunction<InputT, OutputT>) fn);
  }

  /**
   * Returns a type descriptor for the output of the given {@link ProcessFunction}, subject to Java
   * type erasure: may contain unresolved type variables if the type was erased.
   */
  public static <InputT, OutputT> TypeDescriptor<OutputT> outputOf(
      ProcessFunction<InputT, OutputT> fn) {
    return extractFromTypeParameters(
        fn,
        ProcessFunction.class,
        new TypeVariableExtractor<ProcessFunction<InputT, OutputT>, OutputT>() {});
  }

  /** Binary compatibility adapter for {@link #outputOf(ProcessFunction)}. */
  public static <InputT, OutputT> TypeDescriptor<OutputT> outputOf(
      SerializableFunction<InputT, OutputT> fn) {
    return outputOf((ProcessFunction<InputT, OutputT>) fn);
  }

  /** Like {@link #inputOf(ProcessFunction)} but for {@link Contextful.Fn}. */
  public static <InputT, OutputT> TypeDescriptor<InputT> inputOf(
      Contextful.Fn<InputT, OutputT> fn) {
    return TypeDescriptors.extractFromTypeParameters(
        fn,
        Contextful.Fn.class,
        new TypeDescriptors.TypeVariableExtractor<Contextful.Fn<InputT, OutputT>, InputT>() {});
  }

  /** Like {@link #outputOf(ProcessFunction)} but for {@link Contextful.Fn}. */
  public static <InputT, OutputT> TypeDescriptor<OutputT> outputOf(
      Contextful.Fn<InputT, OutputT> fn) {
    return TypeDescriptors.extractFromTypeParameters(
        fn,
        Contextful.Fn.class,
        new TypeDescriptors.TypeVariableExtractor<Contextful.Fn<InputT, OutputT>, OutputT>() {});
  }
}
