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

import org.apache.paimon.predicate.Predicate;
import org.apache.paimon.predicate.PredicateBuilder;
import org.apache.paimon.types.DataType;
import org.apache.paimon.types.RowType;

import org.apache.spark.sql.sources.AlwaysFalse;
import org.apache.spark.sql.sources.AlwaysTrue;
import org.apache.spark.sql.sources.And;
import org.apache.spark.sql.sources.EqualNullSafe;
import org.apache.spark.sql.sources.EqualTo;
import org.apache.spark.sql.sources.Filter;
import org.apache.spark.sql.sources.GreaterThan;
import org.apache.spark.sql.sources.GreaterThanOrEqual;
import org.apache.spark.sql.sources.In;
import org.apache.spark.sql.sources.IsNotNull;
import org.apache.spark.sql.sources.IsNull;
import org.apache.spark.sql.sources.LessThan;
import org.apache.spark.sql.sources.LessThanOrEqual;
import org.apache.spark.sql.sources.Not;
import org.apache.spark.sql.sources.Or;
import org.apache.spark.sql.sources.StringContains;
import org.apache.spark.sql.sources.StringEndsWith;
import org.apache.spark.sql.sources.StringStartsWith;

import javax.annotation.Nullable;

import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import static org.apache.paimon.predicate.PredicateBuilder.convertJavaObject;

/** Conversion from {@link Filter} to {@link Predicate}, remove it when Spark 3.2 is dropped. */
public class SparkFilterConverter {

    public static final List<String> SUPPORT_FILTERS =
            Arrays.asList(
                    "AlwaysTrue",
                    "AlwaysFalse",
                    "EqualTo",
                    "EqualNullSafe",
                    "GreaterThan",
                    "GreaterThanOrEqual",
                    "LessThan",
                    "LessThanOrEqual",
                    "In",
                    "IsNull",
                    "IsNotNull",
                    "And",
                    "Or",
                    "Not",
                    "StringStartsWith",
                    "StringEndsWith",
                    "StringContains");

    private final RowType rowType;
    private final PredicateBuilder builder;

    public SparkFilterConverter(RowType rowType) {
        this.rowType = rowType;
        this.builder = new PredicateBuilder(rowType);
    }

    @Nullable
    public Predicate convertIgnoreFailure(Filter filter) {
        return convert(filter, true);
    }

    @Nullable
    public Predicate convert(Filter filter, boolean ignoreFailure) {
        try {
            return convert(filter);
        } catch (Exception e) {
            if (ignoreFailure) {
                return null;
            } else {
                throw e;
            }
        }
    }

    public Predicate convert(Filter filter) {
        if (filter instanceof AlwaysTrue) {
            return PredicateBuilder.alwaysTrue();
        } else if (filter instanceof AlwaysFalse) {
            return PredicateBuilder.alwaysFalse();
        } else if (filter instanceof EqualTo) {
            EqualTo eq = (EqualTo) filter;
            int index = fieldIndex(eq.attribute());
            if (isNaN(eq.value())) {
                return builder.isNaN(index);
            }
            Object literal = convertLiteral(index, eq.value());
            return builder.equal(index, literal);
        } else if (filter instanceof EqualNullSafe) {
            EqualNullSafe eq = (EqualNullSafe) filter;
            int index = fieldIndex(eq.attribute());
            if (eq.value() == null) {
                return builder.isNull(index);
            } else {
                Object literal = convertLiteral(index, eq.value());
                return builder.equal(index, literal);
            }
        } else if (filter instanceof GreaterThan) {
            GreaterThan gt = (GreaterThan) filter;
            int index = fieldIndex(gt.attribute());
            Object literal = convertLiteral(index, gt.value());
            return builder.greaterThan(index, literal);
        } else if (filter instanceof GreaterThanOrEqual) {
            GreaterThanOrEqual gt = (GreaterThanOrEqual) filter;
            int index = fieldIndex(gt.attribute());
            Object literal = convertLiteral(index, gt.value());
            return builder.greaterOrEqual(index, literal);
        } else if (filter instanceof LessThan) {
            LessThan lt = (LessThan) filter;
            int index = fieldIndex(lt.attribute());
            Object literal = convertLiteral(index, lt.value());
            return builder.lessThan(index, literal);
        } else if (filter instanceof LessThanOrEqual) {
            LessThanOrEqual lt = (LessThanOrEqual) filter;
            int index = fieldIndex(lt.attribute());
            Object literal = convertLiteral(index, lt.value());
            return builder.lessOrEqual(index, literal);
        } else if (filter instanceof In) {
            In in = (In) filter;
            int index = fieldIndex(in.attribute());
            return builder.in(
                    index,
                    Arrays.stream(in.values())
                            .map(v -> convertLiteral(index, v))
                            .collect(Collectors.toList()));
        } else if (filter instanceof IsNull) {
            return builder.isNull(fieldIndex(((IsNull) filter).attribute()));
        } else if (filter instanceof IsNotNull) {
            return builder.isNotNull(fieldIndex(((IsNotNull) filter).attribute()));
        } else if (filter instanceof And) {
            And and = (And) filter;
            return PredicateBuilder.and(convert(and.left()), convert(and.right()));
        } else if (filter instanceof Or) {
            Or or = (Or) filter;
            return PredicateBuilder.or(convert(or.left()), convert(or.right()));
        } else if (filter instanceof Not) {
            Not not = (Not) filter;
            Optional<Predicate> negate = convert(not.child()).negate();
            if (negate.isPresent()) {
                return negate.get();
            }
        } else if (filter instanceof StringStartsWith) {
            StringStartsWith startsWith = (StringStartsWith) filter;
            int index = fieldIndex(startsWith.attribute());
            Object literal = convertLiteral(index, startsWith.value());
            return builder.startsWith(index, literal);
        } else if (filter instanceof StringEndsWith) {
            StringEndsWith endsWith = (StringEndsWith) filter;
            int index = fieldIndex(endsWith.attribute());
            Object literal = convertLiteral(index, endsWith.value());
            return builder.endsWith(index, literal);
        } else if (filter instanceof StringContains) {
            StringContains contains = (StringContains) filter;
            int index = fieldIndex(contains.attribute());
            Object literal = convertLiteral(index, contains.value());
            return builder.contains(index, literal);
        }

        throw new UnsupportedOperationException(
                filter + " is unsupported. Support Filters: " + SUPPORT_FILTERS);
    }

    private static boolean isNaN(Object value) {
        if (value instanceof Float) {
            return Float.isNaN((Float) value);
        }
        if (value instanceof Double) {
            return Double.isNaN((Double) value);
        }
        return false;
    }

    public Object convertLiteral(String field, Object value) {
        return convertLiteral(fieldIndex(field), value);
    }

    public String convertString(String field, Object value) {
        Object literal = convertLiteral(field, value);
        return literal == null ? null : literal.toString();
    }

    private int fieldIndex(String field) {
        int index = rowType.getFieldIndex(field);
        // TODO: support nested field
        if (index == -1) {
            throw new UnsupportedOperationException(
                    String.format("Nested field '%s' is unsupported.", field));
        }
        return index;
    }

    private Object convertLiteral(int index, Object value) {
        DataType type = rowType.getTypeAt(index);
        return convertJavaObject(type, value);
    }
}
