/*
 * 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.flink.connector.mongodb.table;

import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.TableException;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.api.bridge.java.internal.StreamTableEnvironmentImpl;
import org.apache.flink.table.catalog.CatalogManager;
import org.apache.flink.table.catalog.FunctionCatalog;
import org.apache.flink.table.catalog.ResolvedSchema;
import org.apache.flink.table.expressions.Expression;
import org.apache.flink.table.expressions.ResolvedExpression;
import org.apache.flink.table.expressions.resolver.ExpressionResolver;
import org.apache.flink.table.planner.calcite.FlinkContext;
import org.apache.flink.table.planner.calcite.FlinkTypeFactory;
import org.apache.flink.table.planner.calcite.FlinkTypeSystem;
import org.apache.flink.table.planner.delegation.PlannerBase;
import org.apache.flink.table.planner.expressions.RexNodeExpression;
import org.apache.flink.table.planner.plan.utils.FlinkRexUtil;
import org.apache.flink.table.planner.plan.utils.RexNodeToExpressionConverter;
import org.apache.flink.table.types.logical.RowType;

import com.mongodb.client.model.Filters;
import org.apache.calcite.plan.RelOptUtil;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexNode;
import org.bson.BsonBoolean;
import org.bson.BsonDateTime;
import org.bson.BsonDecimal128;
import org.bson.BsonDocument;
import org.bson.BsonInt32;
import org.bson.BsonNull;
import org.bson.BsonString;
import org.bson.conversions.Bson;
import org.bson.types.Decimal128;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.math.BigDecimal;
import java.sql.Timestamp;
import java.time.ZoneId;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.TimeZone;
import java.util.stream.Collectors;

import scala.Option;

import static org.apache.flink.connector.mongodb.table.MongoDynamicTableSource.parseFilter;
import static org.assertj.core.api.Assertions.assertThat;

/** Test for {@link MongoFilterPushDownVisitor}. */
class MongoFilterPushDownVisitorTest {

    private static final String INPUT_TABLE = "mongo_source";

    private static final BsonDocument EMPTY_FILTER = Filters.empty().toBsonDocument();

    private static StreamExecutionEnvironment env;
    private static StreamTableEnvironment tEnv;

    private final ClassLoader classLoader = Thread.currentThread().getContextClassLoader();

    @BeforeEach
    void before() {
        env = StreamExecutionEnvironment.getExecutionEnvironment();
        tEnv = StreamTableEnvironment.create(env);
        tEnv.getConfig().setLocalTimeZone(ZoneId.of("UTC"));

        // Create table in Flink, this can be reused across test cases
        tEnv.executeSql(
                "CREATE TABLE "
                        + INPUT_TABLE
                        + "("
                        + "id INTEGER,"
                        + "description VARCHAR(200),"
                        + "boolean_col BOOLEAN,"
                        + "timestamp_col TIMESTAMP(0),"
                        + "timestamp3_col TIMESTAMP(3),"
                        + "double_col DOUBLE,"
                        + "decimal_col DECIMAL(10, 4)"
                        + ") WITH ("
                        + "  'connector'='mongodb',"
                        + "  'uri'='mongodb://127.0.0.1:27017',"
                        + "  'database'='test_db',"
                        + "  'collection'='test_coll'"
                        + ")");
    }

    @Test
    void testSimpleExpressionPrimitiveType() {
        ResolvedSchema schema = tEnv.sqlQuery("SELECT * FROM " + INPUT_TABLE).getResolvedSchema();
        Arrays.asList(
                        new Object[] {"id = 6", Filters.eq("id", new BsonInt32(6))},
                        new Object[] {"id >= 6", Filters.gte("id", new BsonInt32(6))},
                        new Object[] {"id > 6", Filters.gt("id", new BsonInt32(6))},
                        new Object[] {"id < 6", Filters.lt("id", new BsonInt32(6))},
                        new Object[] {"id <= 5", Filters.lte("id", 5)},
                        new Object[] {
                            "description = 'Halo'",
                            Filters.eq("description", new BsonString("Halo"))
                        },
                        new Object[] {
                            "boolean_col = true", Filters.eq("boolean_col", new BsonBoolean(true))
                        },
                        new Object[] {
                            "boolean_col = false", Filters.eq("boolean_col", new BsonBoolean(false))
                        },
                        new Object[] {
                            "double_col > 0.5",
                            Filters.gt(
                                    "double_col",
                                    new BsonDecimal128(new Decimal128(new BigDecimal("0.5"))))
                        },
                        new Object[] {
                            "decimal_col <= -0.3",
                            Filters.lte(
                                    "decimal_col",
                                    new BsonDecimal128(new Decimal128(new BigDecimal("-0.3"))))
                        })
                .forEach(
                        inputs ->
                                assertGeneratedFilter(
                                        (String) inputs[0],
                                        schema,
                                        ((Bson) inputs[1]).toBsonDocument()));
    }

    @Test
    void testComplexExpressionDatetime() {
        ResolvedSchema schema = tEnv.sqlQuery("SELECT * FROM " + INPUT_TABLE).getResolvedSchema();
        assertGeneratedFilter(
                "id = 6 AND timestamp_col = TIMESTAMP '2022-01-01 07:00:01'",
                schema,
                Filters.and(
                                Filters.eq("id", new BsonInt32(6)),
                                Filters.eq(
                                        "timestamp_col",
                                        new BsonDateTime(
                                                Timestamp.valueOf("2022-01-01 07:00:01")
                                                        .getTime())))
                        .toBsonDocument());

        assertGeneratedFilter(
                "timestamp3_col = TIMESTAMP '2022-01-01 07:00:01.333' OR description = 'Halo'",
                schema,
                Filters.or(
                                Filters.eq(
                                        "timestamp3_col",
                                        new BsonDateTime(
                                                Timestamp.valueOf("2022-01-01 07:00:01.333")
                                                        .getTime())),
                                Filters.eq("description", new BsonString("Halo")))
                        .toBsonDocument());
    }

    @Test
    void testExpressionWithNull() {
        ResolvedSchema schema = tEnv.sqlQuery("SELECT * FROM " + INPUT_TABLE).getResolvedSchema();

        assertGeneratedFilter(
                "id = NULL AND decimal_col <= 0.6",
                schema,
                Filters.and(
                                Filters.eq("id", BsonNull.VALUE),
                                Filters.lte(
                                        "decimal_col",
                                        new BsonDecimal128(new Decimal128(new BigDecimal("0.6")))))
                        .toBsonDocument());

        assertGeneratedFilter(
                "id = 6 OR description = NULL",
                schema,
                Filters.or(
                                Filters.eq("id", new BsonInt32(6)),
                                Filters.eq("description", BsonNull.VALUE))
                        .toBsonDocument());
    }

    @Test
    void testExpressionIsNull() {
        ResolvedSchema schema = tEnv.sqlQuery("SELECT * FROM " + INPUT_TABLE).getResolvedSchema();

        assertGeneratedFilter(
                "id IS NULL AND decimal_col <= 0.6",
                schema,
                Filters.and(
                                Filters.eq("id", BsonNull.VALUE),
                                Filters.lte(
                                        "decimal_col",
                                        new BsonDecimal128(new Decimal128(new BigDecimal("0.6")))))
                        .toBsonDocument());

        assertGeneratedFilter(
                "id = 6 OR description IS NOT NULL",
                schema,
                Filters.or(
                                Filters.eq("id", new BsonInt32(6)),
                                Filters.ne("description", BsonNull.VALUE))
                        .toBsonDocument());
    }

    @Test
    void testExpressionCannotBePushedDown() {
        ResolvedSchema schema = tEnv.sqlQuery("SELECT * FROM " + INPUT_TABLE).getResolvedSchema();

        // unsupported operators
        assertGeneratedFilter("description LIKE '_bcd%'", schema, EMPTY_FILTER);

        // nested complex expressions
        assertGeneratedFilter("double_col = decimal_col", schema, EMPTY_FILTER);
        assertGeneratedFilter("boolean_col = (decimal_col > 2.0)", schema, EMPTY_FILTER);

        // partial push down
        assertGeneratedFilter(
                "id IS NULL AND description LIKE '_bcd%'",
                schema, Filters.eq("id", BsonNull.VALUE).toBsonDocument());

        // sub filter cannot be pushed down
        assertGeneratedFilter("id IS NOT NULL OR double_col = decimal_col", schema, EMPTY_FILTER);
    }

    private void assertGeneratedFilter(
            String inputExpr, ResolvedSchema schema, BsonDocument expected) {
        List<ResolvedExpression> filters = resolveSQLFilterToExpression(inputExpr, schema);

        List<Bson> mongoFilters = new ArrayList<>();
        for (ResolvedExpression filter : filters) {
            BsonDocument simpleFilter = parseFilter(filter);
            if (!simpleFilter.isEmpty()) {
                mongoFilters.add(simpleFilter);
            }
        }

        BsonDocument actual = EMPTY_FILTER;
        if (!mongoFilters.isEmpty()) {
            Bson mergedFilter =
                    mongoFilters.size() == 1 ? mongoFilters.get(0) : Filters.and(mongoFilters);
            actual = mergedFilter.toBsonDocument();
        }

        assertThat(actual).isEqualTo(expected);
    }

    /**
     * Resolve a SQL filter expression against a Schema, this method makes use of some
     * implementation details of Flink.
     */
    private List<ResolvedExpression> resolveSQLFilterToExpression(
            String sqlExp, ResolvedSchema schema) {
        StreamTableEnvironmentImpl tbImpl = (StreamTableEnvironmentImpl) tEnv;

        FlinkContext ctx = ((PlannerBase) tbImpl.getPlanner()).getFlinkContext();
        CatalogManager catMan = tbImpl.getCatalogManager();
        FunctionCatalog funCat = ctx.getFunctionCatalog();
        RowType sourceType = (RowType) schema.toSourceRowDataType().getLogicalType();

        FlinkTypeFactory typeFactory = new FlinkTypeFactory(classLoader, FlinkTypeSystem.INSTANCE);
        RexBuilder rexBuilder = new RexBuilder(typeFactory);
        RexNodeToExpressionConverter converter =
                new RexNodeToExpressionConverter(
                        rexBuilder,
                        sourceType.getFieldNames().toArray(new String[0]),
                        funCat,
                        catMan,
                        TimeZone.getTimeZone(tEnv.getConfig().getLocalTimeZone()));

        RexNodeExpression rexExp =
                (RexNodeExpression) tbImpl.getParser().parseSqlExpression(sqlExp, sourceType, null);

        RexNode cnf = FlinkRexUtil.toCnf(rexBuilder, -1, rexExp.getRexNode());
        // converts the cnf condition to a list of AND conditions
        List<RexNode> conjunctions = RelOptUtil.conjunctions(cnf);

        List<Expression> resolvedExps =
                conjunctions.stream()
                        .map(rex -> rex.accept(converter))
                        .filter(Option::isDefined)
                        .map(Option::get)
                        .collect(Collectors.toList());

        ExpressionResolver resolver =
                ExpressionResolver.resolverFor(
                                tEnv.getConfig(),
                                classLoader,
                                name -> Optional.empty(),
                                funCat.asLookup(
                                        str -> {
                                            throw new TableException(
                                                    "We should not need to lookup any expressions at this point");
                                        }),
                                catMan.getDataTypeFactory(),
                                (sqlExpression, inputRowType, outputType) -> {
                                    throw new TableException(
                                            "SQL expression parsing is not supported at this location.");
                                })
                        .build();

        return resolver.resolve(resolvedExps);
    }
}