/*
 * 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.schemas.transforms.providers;

import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;

import com.google.auto.value.AutoValue;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.File;
import java.io.IOException;
import java.io.Serializable;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.lang.reflect.Type;
import java.math.BigDecimal;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.channels.FileChannel;
import java.nio.channels.ReadableByteChannel;
import java.nio.file.StandardOpenOption;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.annotation.Nullable;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.schemas.AutoValueSchema;
import org.apache.beam.sdk.schemas.FieldValueTypeInformation;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
import org.apache.beam.sdk.schemas.utils.StaticSchemaInference;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams;

public class JavaRowUdf implements Serializable {
  private final Configuration config;
  private final Schema inputSchema;
  private final Schema.FieldType outputType;

  // Transient so we don't have to worry about issues serializing these dynamically created classes.
  // While this is lazily computed, it is always computed on class construction, so any errors
  // should still be caught at construction time, and lazily re-computed before any use.
  @SuppressFBWarnings("SE_TRANSIENT_FIELD_NOT_RESTORED")
  private transient Function<Row, Object> function;

  // Find or implement the inverse of StaticSchemaInference.fieldFromType
  @DefaultSchema(AutoValueSchema.class)
  @AutoValue
  public abstract static class Configuration implements Serializable {
    @SchemaFieldDescription("Source code of a java expression in terms of the schema fields.")
    @Nullable
    public abstract String getExpression();

    @SchemaFieldDescription(
        "Source code of a public class implementing Function<Row, T> for some schema-compatible T.")
    @Nullable
    public abstract String getCallable();

    @SchemaFieldDescription("Path to a jar file implementing the function referenced in name.")
    @Nullable
    public abstract String getPath();

    @SchemaFieldDescription(
        "Fully qualified name of either a class implementing Function<Row, T> (e.g. com.pkg.MyFunction), "
            + "or a method taking a single Row argument (e.g. com.pkg.MyClass::methodName). "
            + "If a method is passed, it must either be static or belong to a class with a public nullary constructor.")
    @Nullable
    public abstract String getName();

    public void validate() {
      checkArgument(
          Strings.isNullOrEmpty(getPath()) || !Strings.isNullOrEmpty(getName()),
          "Specifying a path only allows if a name is provided.");
      int totalArgs =
          (Strings.isNullOrEmpty(getExpression()) ? 0 : 1)
              + (Strings.isNullOrEmpty(getCallable()) ? 0 : 1)
              + (Strings.isNullOrEmpty(getName()) ? 0 : 1);
      checkArgument(
          totalArgs == 1, "Exactly one of expression, callable, or name must be provided.");
    }

    public static Configuration.Builder builder() {
      return new AutoValue_JavaRowUdf_Configuration.Builder();
    }

    @AutoValue.Builder
    public abstract static class Builder {
      public abstract Configuration.Builder setExpression(String expression);

      public abstract Configuration.Builder setCallable(String callable);

      public abstract Configuration.Builder setPath(String path);

      public abstract Configuration.Builder setName(String name);

      public abstract Configuration build();
    }
  }

  public JavaRowUdf(Configuration config, Schema inputSchema)
      throws MalformedURLException, ReflectiveOperationException, StringCompiler.CompileException {
    this.config = config;
    this.inputSchema = inputSchema;
    FunctionAndType functionAndType = createFunction(config, inputSchema);
    this.outputType = functionAndType.outputType;
    this.function = functionAndType.function;
  }

  public Schema.FieldType getOutputType() {
    return outputType;
  }

  public Function<Row, Object> getFunction()
      throws MalformedURLException, ReflectiveOperationException, StringCompiler.CompileException {
    if (function == null) {
      FunctionAndType functionAndType = createFunction(config, inputSchema);
      assert functionAndType.outputType.equals(outputType);
      function = functionAndType.function;
    }
    return function;
  }

  private static class FunctionAndType {
    public final Schema.FieldType outputType;
    public final Function<Row, Object> function;

    public FunctionAndType(Function<Row, Object> function) {
      this(outputOf(function), function);
    }

    public FunctionAndType(Type outputType, Function<Row, Object> function) {
      this(TypeDescriptor.of(outputType), function);
    }

    public FunctionAndType(TypeDescriptor<?> outputType, Function<Row, Object> function) {
      this(
          StaticSchemaInference.fieldFromType(outputType, new EmptyFieldValueTypeSupplier()),
          function);
    }

    public FunctionAndType(Schema.FieldType outputType, Function<Row, Object> function) {
      this.outputType = outputType;
      this.function = function;
    }

    public static <InputT, OutputT> TypeDescriptor<OutputT> outputOf(Function<InputT, OutputT> fn) {
      return TypeDescriptors.extractFromTypeParameters(
          fn,
          Function.class,
          new TypeDescriptors.TypeVariableExtractor<Function<InputT, OutputT>, OutputT>() {});
    }
  }

  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  private static FunctionAndType createFunction(Configuration config, Schema inputSchema)
      throws ReflectiveOperationException, StringCompiler.CompileException, MalformedURLException {
    config.validate();
    if (!Strings.isNullOrEmpty(config.getExpression())) {
      return createFunctionFromExpression(config.getExpression(), inputSchema);
    } else if (!Strings.isNullOrEmpty(config.getCallable())) {
      return createFuctionFromCallable(config.getCallable());
    } else if (!Strings.isNullOrEmpty(config.getName())) {
      return createFunctionFromName(config.getName(), config.getPath());
    } else {
      throw new UnsupportedOperationException(config.toString());
    }
  }

  private static FunctionAndType createFunctionFromExpression(String expression, Schema inputSchema)
      throws StringCompiler.CompileException, ReflectiveOperationException {
    if (inputSchema.hasField(expression)) {
      final int ix = inputSchema.indexOf(expression);
      return new FunctionAndType(
          inputSchema.getField(expression).getType(), (Row row) -> row.getValue(ix));
    } else {
      Map<String, Type> fieldTypes = new HashMap<>();
      for (Schema.Field field : inputSchema.getFields()) {
        if (expression.indexOf(field.getName()) != -1) {
          fieldTypes.put(field.getName(), typeFromFieldType(field.getType()));
        }
      }
      Type type = StringCompiler.guessExpressionType(expression, fieldTypes);
      StringBuilder source = new StringBuilder();
      source.append("import java.util.function.Function;\n");
      source.append("import " + Row.class.getTypeName() + ";\n");
      source.append("public class Eval implements Function<Row, Object> {\n");
      source.append("  public Object apply(Row __row__) {\n");
      for (Map.Entry<String, Type> fieldEntry : fieldTypes.entrySet()) {
        source.append(
            String.format(
                "    %s %s = (%s) __row__.getValue(%s);%n",
                fieldEntry.getValue().getTypeName(),
                fieldEntry.getKey(),
                fieldEntry.getValue().getTypeName(),
                inputSchema.indexOf(fieldEntry.getKey())));
      }
      source.append("    return " + expression + ";\n");
      source.append("  }\n");
      source.append("}\n");
      return new FunctionAndType(
          type, (Function<Row, Object>) StringCompiler.getInstance("Eval", source.toString()));
    }
  }

  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  private static FunctionAndType createFuctionFromCallable(String callable)
      throws StringCompiler.CompileException, ReflectiveOperationException {
    Matcher matcher =
        Pattern.compile("\\bpublic\\s+class\\s+(\\S+)", Pattern.MULTILINE).matcher(callable);
    Preconditions.checkArgument(matcher.find(), "No public class defined in callable source.");
    return new FunctionAndType(
        (Function<Row, Object>) StringCompiler.getInstance(matcher.group(1), callable.toString()));
  }

  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  private static FunctionAndType createFunctionFromName(String name, String path)
      throws ReflectiveOperationException, MalformedURLException {
    if (path != null && !new File(path).exists()) {
      try (ReadableByteChannel inChannel =
          FileSystems.open(FileSystems.matchNewResource(path, false))) {
        File tmpJar = File.createTempFile("map-to-fields-" + name, ".jar");
        try (FileChannel outChannel = FileChannel.open(tmpJar.toPath(), StandardOpenOption.WRITE)) {
          ByteStreams.copy(inChannel, outChannel);
        }
        path = tmpJar.getPath();
      } catch (IOException exn) {
        throw new RuntimeException(exn);
      }
    }
    ClassLoader classLoader =
        path == null
            ? ClassLoader.getSystemClassLoader()
            : new URLClassLoader(
                new URL[] {new URL("file://" + path)}, ClassLoader.getSystemClassLoader());
    String className, methodName = null;
    if (name.indexOf("::") == -1) {
      className = name;
      methodName = null;
    } else {
      String[] parts = name.split("::", 2);
      className = parts[0];
      methodName = parts[1];
    }
    if (methodName == null) {
      return new FunctionAndType(
          (Function<Row, Object>)
              classLoader.loadClass(className).getDeclaredConstructor().newInstance());
    } else {
      Class<?> clazz = classLoader.loadClass(className);
      Method method = clazz.getMethod(methodName, Row.class);
      Object base =
          Modifier.isStatic(method.getModifiers())
              ? null
              : clazz.getDeclaredConstructor().newInstance();
      return new FunctionAndType(
          method.getGenericReturnType(),
          (Row row) -> {
            try {
              return method.invoke(base, row);
            } catch (IllegalAccessException | InvocationTargetException exn) {
              throw new RuntimeException(exn);
            }
          });
    }
  }

  private static class EmptyFieldValueTypeSupplier
      implements org.apache.beam.sdk.schemas.utils.FieldValueTypeSupplier {
    @Override
    public List<FieldValueTypeInformation> get(Class<?> clazz) {
      return Collections.<FieldValueTypeInformation>emptyList();
    }
  }

  private static final Map<Schema.TypeName, Type> NULLABLE_PRIMITIVES =
      ImmutableMap.<Schema.TypeName, Type>builder()
          .put(Schema.TypeName.BYTE, Byte.class)
          .put(Schema.TypeName.INT16, Short.class)
          .put(Schema.TypeName.INT32, Integer.class)
          .put(Schema.TypeName.INT64, Long.class)
          .put(Schema.TypeName.FLOAT, Float.class)
          .put(Schema.TypeName.DOUBLE, Double.class)
          .put(Schema.TypeName.BOOLEAN, Boolean.class)
          .put(Schema.TypeName.BYTES, byte[].class)
          .put(Schema.TypeName.STRING, String.class)
          .put(Schema.TypeName.DECIMAL, BigDecimal.class)
          .build();

  private static final Map<Schema.TypeName, Type> NON_NULLABLE_PRIMITIVES =
      ImmutableMap.<Schema.TypeName, Type>builder()
          .put(Schema.TypeName.BYTE, byte.class)
          .put(Schema.TypeName.INT16, short.class)
          .put(Schema.TypeName.INT32, int.class)
          .put(Schema.TypeName.INT64, long.class)
          .put(Schema.TypeName.FLOAT, float.class)
          .put(Schema.TypeName.DOUBLE, double.class)
          .put(Schema.TypeName.BOOLEAN, boolean.class)
          .put(Schema.TypeName.BYTES, byte[].class)
          .put(Schema.TypeName.STRING, String.class)
          .put(Schema.TypeName.DECIMAL, BigDecimal.class)
          .build();

  private static Type typeFromFieldType(Schema.FieldType fieldType) {
    Map<Schema.TypeName, Type> primitivesMap =
        fieldType.getNullable() ? NULLABLE_PRIMITIVES : NON_NULLABLE_PRIMITIVES;
    if (primitivesMap.containsKey(fieldType.getTypeName())) {
      return primitivesMap.get(fieldType.getTypeName());
    } else if (fieldType.getRowSchema() != null) {
      return Row.class;
    } else {
      throw new UnsupportedOperationException(fieldType.toString());
    }
  }
}
