/*
 * 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.table.planner.plan.rules.logical;

import org.apache.flink.table.api.TableConfig;
import org.apache.flink.table.planner.plan.nodes.FlinkConventions;
import org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc;
import org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalTableSourceScan;
import org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalWatermarkAssigner;
import org.apache.flink.table.planner.plan.optimize.program.FlinkChainedProgram;
import org.apache.flink.table.planner.plan.optimize.program.FlinkHepRuleSetProgramBuilder;
import org.apache.flink.table.planner.plan.optimize.program.FlinkVolcanoProgramBuilder;
import org.apache.flink.table.planner.plan.optimize.program.HEP_RULES_EXECUTION_TYPE;
import org.apache.flink.table.planner.plan.optimize.program.StreamOptimizeContext;
import org.apache.flink.table.planner.runtime.utils.JavaUserDefinedScalarFunctions;
import org.apache.flink.table.planner.runtime.utils.JavaUserDefinedScalarFunctions.JavaFunc5;
import org.apache.flink.table.planner.utils.StreamTableTestUtil;
import org.apache.flink.table.planner.utils.TableTestBase;

import org.apache.calcite.plan.Convention;
import org.apache.calcite.plan.hep.HepMatchOrder;
import org.apache.calcite.rel.rules.CoreRules;
import org.apache.calcite.tools.RuleSets;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.time.Duration;

import static org.apache.flink.table.api.config.ExecutionConfigOptions.TABLE_EXEC_SOURCE_IDLE_TIMEOUT;

/**
 * Test rule {@link PushWatermarkIntoTableSourceScanAcrossCalcRule} and {@link
 * PushWatermarkIntoTableSourceScanRule}.
 */
class PushWatermarkIntoTableSourceScanRuleTest extends TableTestBase {
    private final StreamTableTestUtil util = streamTestUtil(TableConfig.getDefault());

    @BeforeEach
    public void setup() {
        FlinkChainedProgram<StreamOptimizeContext> program = new FlinkChainedProgram<>();
        program.addLast(
                "Converters",
                FlinkVolcanoProgramBuilder.<StreamOptimizeContext>newBuilder()
                        .add(
                                RuleSets.ofList(
                                        CoreRules.PROJECT_TO_CALC,
                                        CoreRules.FILTER_TO_CALC,
                                        FlinkLogicalCalc.CONVERTER(),
                                        FlinkLogicalTableSourceScan.CONVERTER(),
                                        FlinkLogicalWatermarkAssigner.CONVERTER()))
                        .setRequiredOutputTraits(new Convention[] {FlinkConventions.LOGICAL()})
                        .build());
        program.addLast(
                "PushWatermarkIntoTableSourceScanRule",
                FlinkHepRuleSetProgramBuilder.<StreamOptimizeContext>newBuilder()
                        .setHepRulesExecutionType(HEP_RULES_EXECUTION_TYPE.RULE_SEQUENCE())
                        .setHepMatchOrder(HepMatchOrder.BOTTOM_UP)
                        .add(
                                RuleSets.ofList(
                                        PushWatermarkIntoTableSourceScanRule.INSTANCE,
                                        PushWatermarkIntoTableSourceScanAcrossCalcRule.INSTANCE))
                        .build());
        util.replaceStreamProgram(program);
    }

    @Test
    void testSimpleWatermark() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("select a, c from MyTable");
    }

    @Test
    void testWatermarkOnComputedColumn() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  d AS c + INTERVAL '5' SECOND,\n"
                        + "  WATERMARK FOR d AS d - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + " 'connector' = 'values',\n"
                        + " 'enable-watermark-push-down' = 'true',\n"
                        + " 'bounded' = 'false',\n"
                        + " 'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("SELECT * from MyTable");
    }

    @Test
    void testWatermarkOnComputedColumnWithQuery() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3) NOT NULL,\n"
                        + "  d AS c + INTERVAL '5' SECOND,\n"
                        + "  WATERMARK FOR d AS d - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan(
                "SELECT a, b FROM MyTable WHERE d > TO_TIMESTAMP('2020-10-09 12:12:12')");
    }

    @Test
    void testWatermarkOnComputedColumnWithMultipleInputs() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a STRING,\n"
                        + "  b STRING,\n"
                        + "  c as TO_TIMESTAMP(a, b),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("SELECT * FROM MyTable");
    }

    @Test
    void testWatermarkOnRow() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c ROW<name STRING, d TIMESTAMP(3)>,"
                        + "  e AS c.d,"
                        + "  WATERMARK FOR e AS e - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("SELECT * FROM MyTable");
    }

    @Test
    void testWatermarkOnNestedRow() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c ROW<name STRING, d row<e STRING, f TIMESTAMP(3)>>,"
                        + "  g as c.d.f,"
                        + "  WATERMARK for g as g - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("SELECT * FROM MyTable");
    }

    @Test
    void testWatermarkWithMultiInputUdf() {
        JavaFunc5.closeCalled = false;
        JavaFunc5.openCalled = false;
        util.addTemporarySystemFunction("func", new JavaFunc5());

        String ddl =
                "CREATE TABLE MyTable(\n"
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  d AS func(c, a),\n"
                        + "  WATERMARK FOR d AS func(func(d, a), a)\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("SELECT * FROM MyTable");
    }

    @Test
    void testWatermarkOnMetadata() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  `a` INT,\n"
                        + "  `b` BIGINT,\n"
                        + "  `c` TIMESTAMP(3),\n"
                        + "  `metadata` BIGINT METADATA FROM 'metadata_2' VIRTUAL,\n"
                        + "  `computed` AS `metadata` + `b`,\n"
                        + "  WATERMARK for `c` as c - CAST(`metadata` + `computed` AS INTERVAL SECOND)\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'readable-metadata' = 'metadata_1:STRING,metadata_2:INT',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("SELECT * FROM MyTable");
    }

    @Test
    void testWatermarkWithIdleSource() {
        util.tableEnv().getConfig().set(TABLE_EXEC_SOURCE_IDLE_TIMEOUT, Duration.ofMillis(1000));
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("select a, c from MyTable");
    }

    @Test
    void testWatermarkWithPythonFunctionInComputedColumn() {
        util.tableEnv()
                .createTemporaryFunction(
                        "parse_ts",
                        new JavaUserDefinedScalarFunctions.PythonTimestampScalarFunction());
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b AS parse_ts(a),\n"
                        + "  WATERMARK FOR b AS b\n"
                        + ") WITH (\n"
                        + " 'connector' = 'values',\n"
                        + " 'enable-watermark-push-down' = 'true',\n"
                        + " 'bounded' = 'false',\n"
                        + " 'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("SELECT * FROM MyTable");
    }

    @Test
    void testWatermarkEmitStrategyWithOptions() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'scan.watermark.emit.strategy' = 'on-event',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("select a, c from MyTable");
    }

    @Test
    void testWatermarkEmitStrategyWithHint() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan(
                "select a, c from MyTable /*+ OPTIONS("
                        + "'scan.watermark.emit.strategy'='on-event') */");
    }

    @Test
    void testWatermarkAlignmentWithOptions() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'scan.watermark.alignment.group' = 'group1',\n"
                        + "  'scan.watermark.alignment.max-drift' = '1min',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("select a, c from MyTable");
    }

    @Test
    void testWatermarkAlignmentWithHint() {
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan(
                "select a, c from MyTable /*+ OPTIONS("
                        + "'scan.watermark.alignment.group'='group1', "
                        + "'scan.watermark.alignment.max-drift'='1min') */");
    }

    @Test
    void testIdleSourceWithOptions() {
        util.tableEnv().getConfig().set(TABLE_EXEC_SOURCE_IDLE_TIMEOUT, Duration.ofMillis(1000));
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true',\n"
                        + "  'scan.watermark.idle-timeout' = '60s'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan("select a, c from MyTable");
    }

    @Test
    void testIdleSourceWithHint() {
        util.tableEnv().getConfig().set(TABLE_EXEC_SOURCE_IDLE_TIMEOUT, Duration.ofMillis(1000));
        String ddl =
                "CREATE TABLE MyTable("
                        + "  a INT,\n"
                        + "  b BIGINT,\n"
                        + "  c TIMESTAMP(3),\n"
                        + "  WATERMARK FOR c AS c - INTERVAL '5' SECOND\n"
                        + ") WITH (\n"
                        + "  'connector' = 'values',\n"
                        + "  'enable-watermark-push-down' = 'true',\n"
                        + "  'bounded' = 'false',\n"
                        + "  'disable-lookup' = 'true'"
                        + ")";
        util.tableEnv().executeSql(ddl);
        util.verifyRelPlan(
                "select a, c from MyTable /*+ OPTIONS('scan.watermark.idle-timeout' = '60s')*/");
    }
}
