/*
 * 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.iceberg.spark;

import static org.apache.iceberg.expressions.Expressions.and;
import static org.apache.iceberg.expressions.Expressions.bucket;
import static org.apache.iceberg.expressions.Expressions.day;
import static org.apache.iceberg.expressions.Expressions.equal;
import static org.apache.iceberg.expressions.Expressions.greaterThan;
import static org.apache.iceberg.expressions.Expressions.greaterThanOrEqual;
import static org.apache.iceberg.expressions.Expressions.hour;
import static org.apache.iceberg.expressions.Expressions.in;
import static org.apache.iceberg.expressions.Expressions.isNaN;
import static org.apache.iceberg.expressions.Expressions.isNull;
import static org.apache.iceberg.expressions.Expressions.lessThan;
import static org.apache.iceberg.expressions.Expressions.lessThanOrEqual;
import static org.apache.iceberg.expressions.Expressions.month;
import static org.apache.iceberg.expressions.Expressions.not;
import static org.apache.iceberg.expressions.Expressions.notEqual;
import static org.apache.iceberg.expressions.Expressions.notIn;
import static org.apache.iceberg.expressions.Expressions.notNaN;
import static org.apache.iceberg.expressions.Expressions.notNull;
import static org.apache.iceberg.expressions.Expressions.or;
import static org.apache.iceberg.expressions.Expressions.startsWith;
import static org.apache.iceberg.expressions.Expressions.truncate;
import static org.apache.iceberg.expressions.Expressions.year;

import java.util.Arrays;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import org.apache.iceberg.expressions.Expression;
import org.apache.iceberg.expressions.Expression.Operation;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.expressions.UnboundPredicate;
import org.apache.iceberg.expressions.UnboundTerm;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet;
import org.apache.iceberg.util.NaNUtil;
import org.apache.iceberg.util.Pair;
import org.apache.spark.sql.connector.expressions.Literal;
import org.apache.spark.sql.connector.expressions.NamedReference;
import org.apache.spark.sql.connector.expressions.UserDefinedScalarFunc;
import org.apache.spark.sql.connector.expressions.filter.And;
import org.apache.spark.sql.connector.expressions.filter.Not;
import org.apache.spark.sql.connector.expressions.filter.Or;
import org.apache.spark.sql.connector.expressions.filter.Predicate;
import org.apache.spark.unsafe.types.UTF8String;

public class SparkV2Filters {

  public static final Set<String> SUPPORTED_FUNCTIONS =
      ImmutableSet.of("years", "months", "days", "hours", "bucket", "truncate");

  private static final String TRUE = "ALWAYS_TRUE";
  private static final String FALSE = "ALWAYS_FALSE";
  private static final String EQ = "=";
  private static final String EQ_NULL_SAFE = "<=>";
  private static final String NOT_EQ = "<>";
  private static final String GT = ">";
  private static final String GT_EQ = ">=";
  private static final String LT = "<";
  private static final String LT_EQ = "<=";
  private static final String IN = "IN";
  private static final String IS_NULL = "IS_NULL";
  private static final String NOT_NULL = "IS_NOT_NULL";
  private static final String AND = "AND";
  private static final String OR = "OR";
  private static final String NOT = "NOT";
  private static final String STARTS_WITH = "STARTS_WITH";

  private static final Map<String, Operation> FILTERS =
      ImmutableMap.<String, Operation>builder()
          .put(TRUE, Operation.TRUE)
          .put(FALSE, Operation.FALSE)
          .put(EQ, Operation.EQ)
          .put(EQ_NULL_SAFE, Operation.EQ)
          .put(NOT_EQ, Operation.NOT_EQ)
          .put(GT, Operation.GT)
          .put(GT_EQ, Operation.GT_EQ)
          .put(LT, Operation.LT)
          .put(LT_EQ, Operation.LT_EQ)
          .put(IN, Operation.IN)
          .put(IS_NULL, Operation.IS_NULL)
          .put(NOT_NULL, Operation.NOT_NULL)
          .put(AND, Operation.AND)
          .put(OR, Operation.OR)
          .put(NOT, Operation.NOT)
          .put(STARTS_WITH, Operation.STARTS_WITH)
          .buildOrThrow();

  private SparkV2Filters() {}

  public static Expression convert(Predicate[] predicates) {
    Expression expression = Expressions.alwaysTrue();
    for (Predicate predicate : predicates) {
      Expression converted = convert(predicate);
      Preconditions.checkArgument(
          converted != null, "Cannot convert Spark predicate to Iceberg expression: %s", predicate);
      expression = Expressions.and(expression, converted);
    }

    return expression;
  }

  @SuppressWarnings({"checkstyle:CyclomaticComplexity", "checkstyle:MethodLength"})
  public static Expression convert(Predicate predicate) {
    Operation op = FILTERS.get(predicate.name());
    if (op != null) {
      switch (op) {
        case TRUE:
          return Expressions.alwaysTrue();

        case FALSE:
          return Expressions.alwaysFalse();

        case IS_NULL:
          if (canConvertToTerm(child(predicate))) {
            UnboundTerm<Object> term = toTerm(child(predicate));
            return term != null ? isNull(term) : null;
          }

          return null;

        case NOT_NULL:
          if (canConvertToTerm(child(predicate))) {
            UnboundTerm<Object> term = toTerm(child(predicate));
            return term != null ? notNull(term) : null;
          }

          return null;

        case LT:
          if (canConvertToTerm(leftChild(predicate)) && isLiteral(rightChild(predicate))) {
            UnboundTerm<Object> term = toTerm(leftChild(predicate));
            return term != null ? lessThan(term, convertLiteral(rightChild(predicate))) : null;
          } else if (canConvertToTerm(rightChild(predicate)) && isLiteral(leftChild(predicate))) {
            UnboundTerm<Object> term = toTerm(rightChild(predicate));
            return term != null ? greaterThan(term, convertLiteral(leftChild(predicate))) : null;
          } else {
            return null;
          }

        case LT_EQ:
          if (canConvertToTerm(leftChild(predicate)) && isLiteral(rightChild(predicate))) {
            UnboundTerm<Object> term = toTerm(leftChild(predicate));
            return term != null
                ? lessThanOrEqual(term, convertLiteral(rightChild(predicate)))
                : null;
          } else if (canConvertToTerm(rightChild(predicate)) && isLiteral(leftChild(predicate))) {
            UnboundTerm<Object> term = toTerm(rightChild(predicate));
            return term != null
                ? greaterThanOrEqual(term, convertLiteral(leftChild(predicate)))
                : null;
          } else {
            return null;
          }

        case GT:
          if (canConvertToTerm(leftChild(predicate)) && isLiteral(rightChild(predicate))) {
            UnboundTerm<Object> term = toTerm(leftChild(predicate));
            return term != null ? greaterThan(term, convertLiteral(rightChild(predicate))) : null;
          } else if (canConvertToTerm(rightChild(predicate)) && isLiteral(leftChild(predicate))) {
            UnboundTerm<Object> term = toTerm(rightChild(predicate));
            return term != null ? lessThan(term, convertLiteral(leftChild(predicate))) : null;
          } else {
            return null;
          }

        case GT_EQ:
          if (canConvertToTerm(leftChild(predicate)) && isLiteral(rightChild(predicate))) {
            UnboundTerm<Object> term = toTerm(leftChild(predicate));
            return term != null
                ? greaterThanOrEqual(term, convertLiteral(rightChild(predicate)))
                : null;
          } else if (canConvertToTerm(rightChild(predicate)) && isLiteral(leftChild(predicate))) {
            UnboundTerm<Object> term = toTerm(rightChild(predicate));
            return term != null
                ? lessThanOrEqual(term, convertLiteral(leftChild(predicate)))
                : null;
          } else {
            return null;
          }

        case EQ: // used for both eq and null-safe-eq
          Pair<UnboundTerm<Object>, Object> eqChildren = predicateChildren(predicate);
          if (eqChildren == null) {
            return null;
          }

          if (predicate.name().equals(EQ)) {
            // comparison with null in normal equality is always null. this is probably a mistake.
            Preconditions.checkNotNull(
                eqChildren.second(),
                "Expression is always false (eq is not null-safe): %s",
                predicate);
          }

          return handleEqual(eqChildren.first(), eqChildren.second());

        case NOT_EQ:
          Pair<UnboundTerm<Object>, Object> notEqChildren = predicateChildren(predicate);
          if (notEqChildren == null) {
            return null;
          }

          // comparison with null in normal equality is always null. this is probably a mistake.
          Preconditions.checkNotNull(
              notEqChildren.second(),
              "Expression is always false (notEq is not null-safe): %s",
              predicate);

          return handleNotEqual(notEqChildren.first(), notEqChildren.second());

        case IN:
          if (isSupportedInPredicate(predicate)) {
            UnboundTerm<Object> term = toTerm(childAtIndex(predicate, 0));

            return term != null
                ? in(
                    term,
                    Arrays.stream(predicate.children())
                        .skip(1)
                        .map(val -> convertLiteral(((Literal<?>) val)))
                        .filter(Objects::nonNull)
                        .collect(Collectors.toList()))
                : null;
          } else {
            return null;
          }

        case NOT:
          Not notPredicate = (Not) predicate;
          Predicate childPredicate = notPredicate.child();
          if (childPredicate.name().equals(IN) && isSupportedInPredicate(childPredicate)) {
            UnboundTerm<Object> term = toTerm(childAtIndex(childPredicate, 0));
            if (term == null) {
              return null;
            }

            // infer an extra notNull predicate for Spark NOT IN filters
            // as Iceberg expressions don't follow the 3-value SQL boolean logic
            // col NOT IN (1, 2) in Spark is equal to notNull(col) && notIn(col, 1, 2) in Iceberg
            Expression notIn =
                notIn(
                    term,
                    Arrays.stream(childPredicate.children())
                        .skip(1)
                        .map(val -> convertLiteral(((Literal<?>) val)))
                        .filter(Objects::nonNull)
                        .collect(Collectors.toList()));
            return and(notNull(term), notIn);
          } else if (hasNoInFilter(childPredicate)) {
            Expression child = convert(childPredicate);
            if (child != null) {
              return not(child);
            }
          }
          return null;

        case AND:
          {
            And andPredicate = (And) predicate;
            Expression left = convert(andPredicate.left());
            Expression right = convert(andPredicate.right());
            if (left != null && right != null) {
              return and(left, right);
            }
            return null;
          }

        case OR:
          {
            Or orPredicate = (Or) predicate;
            Expression left = convert(orPredicate.left());
            Expression right = convert(orPredicate.right());
            if (left != null && right != null) {
              return or(left, right);
            }
            return null;
          }

        case STARTS_WITH:
          String colName = SparkUtil.toColumnName(leftChild(predicate));
          return startsWith(colName, convertLiteral(rightChild(predicate)).toString());
      }
    }

    return null;
  }

  private static Pair<UnboundTerm<Object>, Object> predicateChildren(Predicate predicate) {
    if (canConvertToTerm(leftChild(predicate)) && isLiteral(rightChild(predicate))) {
      UnboundTerm<Object> term = toTerm(leftChild(predicate));
      return term != null ? Pair.of(term, convertLiteral(rightChild(predicate))) : null;

    } else if (canConvertToTerm(rightChild(predicate)) && isLiteral(leftChild(predicate))) {
      UnboundTerm<Object> term = toTerm(rightChild(predicate));
      return term != null ? Pair.of(term, convertLiteral(leftChild(predicate))) : null;

    } else {
      return null;
    }
  }

  @SuppressWarnings("unchecked")
  private static <T> T child(Predicate predicate) {
    org.apache.spark.sql.connector.expressions.Expression[] children = predicate.children();
    Preconditions.checkArgument(
        children.length == 1, "Predicate should have one child: %s", predicate);
    return (T) children[0];
  }

  @SuppressWarnings("unchecked")
  private static <T> T leftChild(Predicate predicate) {
    org.apache.spark.sql.connector.expressions.Expression[] children = predicate.children();
    Preconditions.checkArgument(
        children.length == 2, "Predicate should have two children: %s", predicate);
    return (T) children[0];
  }

  @SuppressWarnings("unchecked")
  private static <T> T rightChild(Predicate predicate) {
    org.apache.spark.sql.connector.expressions.Expression[] children = predicate.children();
    Preconditions.checkArgument(
        children.length == 2, "Predicate should have two children: %s", predicate);
    return (T) children[1];
  }

  @SuppressWarnings("unchecked")
  private static <T> T childAtIndex(Predicate predicate, int index) {
    return (T) predicate.children()[index];
  }

  private static boolean canConvertToTerm(
      org.apache.spark.sql.connector.expressions.Expression expr) {
    return isRef(expr) || isSystemFunc(expr);
  }

  private static boolean isRef(org.apache.spark.sql.connector.expressions.Expression expr) {
    return expr instanceof NamedReference;
  }

  private static boolean isSystemFunc(org.apache.spark.sql.connector.expressions.Expression expr) {
    if (expr instanceof UserDefinedScalarFunc) {
      UserDefinedScalarFunc udf = (UserDefinedScalarFunc) expr;
      return udf.canonicalName().startsWith("iceberg")
          && SUPPORTED_FUNCTIONS.contains(udf.name())
          && Arrays.stream(udf.children()).allMatch(child -> isLiteral(child) || isRef(child));
    }

    return false;
  }

  private static boolean isLiteral(org.apache.spark.sql.connector.expressions.Expression expr) {
    return expr instanceof Literal;
  }

  private static Object convertLiteral(Literal<?> literal) {
    if (literal.value() instanceof UTF8String) {
      return ((UTF8String) literal.value()).toString();
    }
    return literal.value();
  }

  private static UnboundPredicate<Object> handleEqual(UnboundTerm<Object> term, Object value) {
    if (value == null) {
      return isNull(term);
    } else if (NaNUtil.isNaN(value)) {
      return isNaN(term);
    } else {
      return equal(term, value);
    }
  }

  private static UnboundPredicate<Object> handleNotEqual(UnboundTerm<Object> term, Object value) {
    if (NaNUtil.isNaN(value)) {
      return notNaN(term);
    } else {
      return notEqual(term, value);
    }
  }

  private static boolean hasNoInFilter(Predicate predicate) {
    Operation op = FILTERS.get(predicate.name());

    if (op != null) {
      switch (op) {
        case AND:
          And andPredicate = (And) predicate;
          return hasNoInFilter(andPredicate.left()) && hasNoInFilter(andPredicate.right());
        case OR:
          Or orPredicate = (Or) predicate;
          return hasNoInFilter(orPredicate.left()) && hasNoInFilter(orPredicate.right());
        case NOT:
          Not notPredicate = (Not) predicate;
          return hasNoInFilter(notPredicate.child());
        case IN:
          return false;
        default:
          return true;
      }
    }

    return false;
  }

  private static boolean isSupportedInPredicate(Predicate predicate) {
    if (!canConvertToTerm(childAtIndex(predicate, 0))) {
      return false;
    } else {
      return Arrays.stream(predicate.children()).skip(1).allMatch(SparkV2Filters::isLiteral);
    }
  }

  /** Should be called after {@link #canConvertToTerm} passed */
  private static <T> UnboundTerm<Object> toTerm(T input) {
    if (input instanceof NamedReference) {
      return Expressions.ref(SparkUtil.toColumnName((NamedReference) input));
    } else if (input instanceof UserDefinedScalarFunc) {
      return udfToTerm((UserDefinedScalarFunc) input);
    } else {
      return null;
    }
  }

  @SuppressWarnings("checkstyle:CyclomaticComplexity")
  private static UnboundTerm<Object> udfToTerm(UserDefinedScalarFunc udf) {
    org.apache.spark.sql.connector.expressions.Expression[] children = udf.children();
    String udfName = udf.name().toLowerCase(Locale.ROOT);
    if (children.length == 1) {
      org.apache.spark.sql.connector.expressions.Expression child = children[0];
      if (isRef(child)) {
        String column = SparkUtil.toColumnName((NamedReference) child);
        switch (udfName) {
          case "years":
            return year(column);
          case "months":
            return month(column);
          case "days":
            return day(column);
          case "hours":
            return hour(column);
        }
      }
    } else if (children.length == 2) {
      if (isLiteral(children[0]) && isRef(children[1])) {
        String column = SparkUtil.toColumnName((NamedReference) children[1]);
        switch (udfName) {
          case "bucket":
            int numBuckets = (Integer) convertLiteral((Literal<?>) children[0]);
            return bucket(column, numBuckets);
          case "truncate":
            int width = (Integer) convertLiteral((Literal<?>) children[0]);
            return truncate(column, width);
        }
      }
    }

    return null;
  }
}
