/*
 * 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.fluss.flink.source;

import org.apache.fluss.client.table.Table;
import org.apache.fluss.client.table.writer.AppendWriter;
import org.apache.fluss.client.table.writer.UpsertWriter;
import org.apache.fluss.exception.InvalidTableException;
import org.apache.fluss.flink.utils.FlinkTestBase;
import org.apache.fluss.metadata.TablePath;

import org.apache.commons.lang3.RandomUtils;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.EnvironmentSettings;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.api.config.ExecutionConfigOptions;
import org.apache.flink.types.Row;
import org.apache.flink.util.CloseableIterator;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import static org.apache.fluss.flink.FlinkConnectorOptions.BOOTSTRAP_SERVERS;
import static org.apache.fluss.flink.source.testutils.FlinkRowAssertionsUtils.assertResultsIgnoreOrder;
import static org.apache.fluss.flink.source.testutils.FlinkRowAssertionsUtils.collectRowsWithTimeout;
import static org.apache.fluss.server.testutils.FlussClusterExtension.BUILTIN_DATABASE;
import static org.apache.fluss.testutils.DataTestUtils.row;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

/** IT case for using flink sql to read fluss table. */
abstract class FlinkTableSourceBatchITCase extends FlinkTestBase {

    static final String CATALOG_NAME = "testcatalog";
    static final String DEFAULT_DB = "defaultdb";
    protected StreamTableEnvironment tEnv;

    @BeforeEach
    void before() {
        StreamExecutionEnvironment execEnv = StreamExecutionEnvironment.getExecutionEnvironment();
        // create table environment
        tEnv = StreamTableEnvironment.create(execEnv, EnvironmentSettings.inBatchMode());
        // crate catalog using sql
        tEnv.executeSql(
                String.format(
                        "create catalog %s with ('type' = 'fluss', '%s' = '%s')",
                        CATALOG_NAME, BOOTSTRAP_SERVERS.key(), bootstrapServers));
        tEnv.executeSql("use catalog " + CATALOG_NAME);

        tEnv.getConfig().set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 4);
        // create database
        tEnv.executeSql("create database " + DEFAULT_DB);
        tEnv.useDatabase(DEFAULT_DB);
    }

    @AfterEach
    void after() {
        tEnv.useDatabase(BUILTIN_DATABASE);
        tEnv.executeSql(String.format("drop database %s cascade", DEFAULT_DB));
    }

    @Test
    void testScanSingleRowFilter() throws Exception {
        String tableName = prepareSourceTable(new String[] {"name", "id"}, null);
        String query = String.format("SELECT * FROM %s WHERE id = 1 AND name = 'name1'", tableName);

        assertThat(tEnv.explainSql(query))
                .contains(
                        String.format(
                                "TableSourceScan(table=[[testcatalog, defaultdb, %s, "
                                        + "filter=[and(=(id, 1), =(name, _UTF-16LE'name1':VARCHAR(2147483647) CHARACTER SET \"UTF-16LE\"))]]], "
                                        + "fields=[id, address, name])",
                                tableName));
        CloseableIterator<Row> collected = tEnv.executeSql(query).collect();
        List<String> expected = Collections.singletonList("+I[1, address1, name1]");
        assertResultsIgnoreOrder(collected, expected, true);
    }

    @Test
    void testScanSingleRowFilter2() throws Exception {
        String tableName = prepareSourceTable(new String[] {"id", "name"}, null);
        String query = String.format("SELECT * FROM %s WHERE id = 1 AND name = 'name1'", tableName);

        assertThat(tEnv.explainSql(query))
                .contains(
                        String.format(
                                "TableSourceScan(table=[[testcatalog, defaultdb, %s, "
                                        + "filter=[and(=(id, 1), =(name, _UTF-16LE'name1':VARCHAR(2147483647) CHARACTER SET \"UTF-16LE\"))]]], "
                                        + "fields=[id, address, name])",
                                tableName));
        CloseableIterator<Row> collected = tEnv.executeSql(query).collect();
        List<String> expected = Collections.singletonList("+I[1, address1, name1]");
        assertResultsIgnoreOrder(collected, expected, true);
    }

    @Test
    void testScanSingleRowFilter3() throws Exception {
        String tableName = prepareSourceTable(new String[] {"id"}, null);
        String query = String.format("SELECT id,name FROM %s WHERE id = 1", tableName);

        assertThat(tEnv.explainSql(query))
                .contains(
                        String.format(
                                "TableSourceScan(table=[[testcatalog, defaultdb, %s, "
                                        + "filter=[=(id, 1)], "
                                        + "project=[id, name]]], fields=[id, name])",
                                tableName));
        CloseableIterator<Row> collected = tEnv.executeSql(query).collect();
        List<String> expected = Collections.singletonList("+I[1, name1]");
        assertResultsIgnoreOrder(collected, expected, true);
    }

    @Test
    void testScanSingleRowFilterOnPartitionedTable() throws Exception {
        String tableName = prepareSourceTable(new String[] {"id", "dt"}, "dt");
        TablePath tablePath = TablePath.of(DEFAULT_DB, tableName);
        Map<Long, String> partitionNameById =
                waitUntilPartitions(FLUSS_CLUSTER_EXTENSION.getZooKeeperClient(), tablePath);
        Iterator<String> partitionIterator =
                partitionNameById.values().stream().sorted().iterator();
        String partition1 = partitionIterator.next();
        String query =
                String.format("SELECT * FROM %s WHERE id = 1 AND dt='%s'", tableName, partition1);

        assertThat(tEnv.explainSql(query))
                .contains(
                        String.format(
                                "TableSourceScan(table=[[testcatalog, defaultdb, %s, "
                                        + "filter=[and(=(id, 1), =(dt, _UTF-16LE'%s':VARCHAR(2147483647) CHARACTER SET \"UTF-16LE\"))]]], "
                                        + "fields=[id, address, name, dt])\n",
                                tableName, partition1));

        CloseableIterator<Row> collected = tEnv.executeSql(query).collect();
        List<String> expected =
                Collections.singletonList(String.format("+I[1, address1, name1, %s]", partition1));
        assertResultsIgnoreOrder(collected, expected, true);
    }

    @Test
    void testFilterOnLookupSource() throws Exception {
        String srcTableName = String.format("test_src_table_%s", RandomUtils.nextInt());
        tEnv.executeSql(
                String.format(
                        "create table %s ("
                                + "  id int not null,"
                                + "  name varchar,"
                                + "  dt varchar,"
                                + "  dim_dt varchar,"
                                + "  primary key (id, dt) NOT ENFORCED) partitioned by (dt)"
                                + " with ("
                                + "  'bucket.num' = '4', "
                                + "  'table.auto-partition.enabled' = 'true',"
                                + "  'table.auto-partition.time-unit' = 'year')",
                        srcTableName));

        String dimTableName = String.format("test_dim_table_%s", RandomUtils.nextInt());
        tEnv.executeSql(
                String.format(
                        "create table %s ("
                                + "  id int not null,"
                                + "  address varchar,"
                                + "  dt varchar,"
                                + "  primary key (id, dt) NOT ENFORCED) partitioned by (dt)"
                                + " with ("
                                + "  'bucket.num' = '4', "
                                + "  'table.auto-partition.enabled' = 'true',"
                                + "  'table.auto-partition.time-unit' = 'year')",
                        dimTableName));

        TablePath srcTablePath = TablePath.of(DEFAULT_DB, srcTableName);
        Map<Long, String> partitionNameById =
                waitUntilPartitions(FLUSS_CLUSTER_EXTENSION.getZooKeeperClient(), srcTablePath);
        // just pick first partition to insert data
        Iterator<String> partitionIterator =
                partitionNameById.values().stream().sorted().iterator();
        String partition1 = partitionIterator.next();

        // prepare src table data
        try (Table srcTable = conn.getTable(srcTablePath)) {
            UpsertWriter upsertWriter = srcTable.newUpsert().createWriter();
            for (int i = 1; i <= 2; i++) {
                Object[] values = new Object[] {i, "name" + i, partition1, partition1};
                upsertWriter.upsert(row(values));
            }
            upsertWriter.flush();
        }

        TablePath dimTablePath = TablePath.of(DEFAULT_DB, dimTableName);
        // prepare dim table data
        try (Table dimTable = conn.getTable(dimTablePath)) {
            UpsertWriter upsertWriter = dimTable.newUpsert().createWriter();
            for (int i = 1; i <= 2; i++) {
                Object[] values = new Object[] {i, "address" + i, partition1};
                upsertWriter.upsert(row(values));
            }
            upsertWriter.flush();
        }

        tEnv.executeSql(
                String.format(
                        "CREATE TEMPORARY VIEW my_view AS "
                                + "SELECT *, proctime() as proc from %s WHERE id = 1 AND dt = '%s'",
                        srcTableName, partition1));

        CloseableIterator<Row> collected =
                tEnv.executeSql(
                                String.format(
                                        "SELECT src.id, src.name, h.id, h.address FROM my_view src "
                                                + " LEFT JOIN %s FOR SYSTEM_TIME AS OF src.proc as h "
                                                + " ON src.id = h.id and src.dim_dt = h.dt and h.dt <> '%s'",
                                        dimTableName, partition1))
                        .collect();
        List<String> expected = Collections.singletonList("+I[1, name1, null, null]");
        assertResultsIgnoreOrder(collected, expected, true);
    }

    @Test
    void testScanSingleRowFilterException() throws Exception {
        String tableName = prepareSourceTable(new String[] {"id", "name"}, null);
        String query = String.format("SELECT * FROM %s WHERE id = 1", tableName);

        // doesn't have all condition for primary key, doesn't support to execute
        assertThatThrownBy(() -> tEnv.explainSql(query))
                .isInstanceOf(UnsupportedOperationException.class)
                .hasMessage(
                        "Currently, Fluss only support queries on table with datalake enabled"
                                + " or point queries on primary key when it's in batch execution mode.");
    }

    @Test
    void testLakeTableQueryOnLakeDisabledTable() throws Exception {
        String tableName = prepareSourceTable(new String[] {"id", "name"}, null);
        assertThatThrownBy(() -> tEnv.executeSql(String.format("SELECT * FROM %s$lake", tableName)))
                .cause()
                .cause()
                .isInstanceOf(UnsupportedOperationException.class)
                .hasMessage(
                        String.format(
                                "Table %s.%s is not datalake enabled.", DEFAULT_DB, tableName));
    }

    @Test
    void testLimitPrimaryTableScan() throws Exception {
        String tableName = prepareSourceTable(new String[] {"id"}, null);
        // normal scan
        String query = String.format("SELECT * FROM %s limit 2", tableName);
        CloseableIterator<Row> iterRows = tEnv.executeSql(query).collect();
        List<String> collected = collectRowsWithTimeout(iterRows, 2);
        List<String> expected =
                Arrays.asList(
                        "+I[1, address1, name1]",
                        "+I[2, address2, name2]",
                        "+I[3, address3, name3]",
                        "+I[4, address4, name4]",
                        "+I[5, address5, name5]");
        assertThat(collected).isSubsetOf(expected);
        assertThat(collected).hasSize(2);

        // limit which is larger than all the data.
        query = String.format("SELECT * FROM %s limit 10", tableName);
        iterRows = tEnv.executeSql(query).collect();
        collected = collectRowsWithTimeout(iterRows, 5);
        assertThat(collected).isSubsetOf(expected);
        assertThat(collected).hasSize(5);

        // projection scan
        query = String.format("SELECT id, name FROM %s limit 3", tableName);
        iterRows = tEnv.executeSql(query).collect();
        collected = collectRowsWithTimeout(iterRows, 3);
        expected =
                Arrays.asList(
                        "+I[1, name1]",
                        "+I[2, name2]",
                        "+I[3, name3]",
                        "+I[4, name4]",
                        "+I[5, name5]");
        assertThat(collected).isSubsetOf(expected);
        assertThat(collected).hasSize(3);

        // limit out of bounds
        assertThatThrownBy(
                        () ->
                                tEnv.executeSql(
                                        String.format(
                                                "SELECT id, name FROM %s limit 10000", tableName)))
                .hasMessageContaining("LIMIT statement doesn't support greater than 2048");
    }

    @Test
    void testLimitLogTableScan() throws Exception {
        String tableName = prepareLogTable();

        // normal scan
        String query = String.format("SELECT * FROM %s limit 2", tableName);
        CloseableIterator<Row> iterRows = tEnv.executeSql(query).collect();
        List<String> collected = collectRowsWithTimeout(iterRows, 2);
        List<String> expected =
                Arrays.asList(
                        "+I[1, address1, name1]",
                        "+I[2, address2, name2]",
                        "+I[3, address3, name3]",
                        "+I[4, address4, name4]",
                        "+I[5, address5, name5]");
        assertThat(collected).isSubsetOf(expected);
        assertThat(collected).hasSize(2);

        // projection scan
        query = String.format("SELECT id, name FROM %s limit 3", tableName);
        iterRows = tEnv.executeSql(query).collect();
        collected = collectRowsWithTimeout(iterRows, 3);
        expected =
                Arrays.asList(
                        "+I[1, name1]",
                        "+I[2, name2]",
                        "+I[3, name3]",
                        "+I[4, name4]",
                        "+I[5, name5]");
        assertThat(collected).isSubsetOf(expected);
        assertThat(collected).hasSize(3);

        // test partition table.
        String partitionTable = preparePartitionedLogTable();
        query = String.format("SELECT id, name FROM %s limit 3", partitionTable);
        iterRows = tEnv.executeSql(query).collect();
        collected = collectRowsWithTimeout(iterRows, 3);
        assertThat(collected).isSubsetOf(expected);
        assertThat(collected).hasSize(3);
    }

    @ParameterizedTest
    @ValueSource(booleans = {true, false})
    void testCountPushDownForPkTable(boolean partitionTable) throws Exception {
        String tableName =
                partitionTable
                        ? prepareSourceTable(new String[] {"id", "dt"}, "dt")
                        : prepareSourceTable(new String[] {"id"}, null);
        // normal scan
        String query = String.format("SELECT COUNT(*) FROM %s", tableName);
        assertThat(tEnv.explainSql(query))
                .contains(
                        "aggregates=[grouping=[], aggFunctions=[Count1AggFunction()]]]], fields=[count1$0]");
        CloseableIterator<Row> iterRows = tEnv.executeSql(query).collect();
        List<String> collected = collectRowsWithTimeout(iterRows, 1);
        List<String> expected = Collections.singletonList("+I[5]");
        assertThat(collected).isEqualTo(expected);

        // test not push down grouping count.
        assertThatThrownBy(
                        () ->
                                tEnv.explainSql(
                                                String.format(
                                                        "SELECT COUNT(*) FROM %s group by id",
                                                        tableName))
                                        .wait())
                .hasMessageContaining(
                        "Currently, Fluss only support queries on table with datalake enabled or point queries on primary key when it's in batch execution mode.");
    }

    @Test
    void testCountPushDownWithWALMode() throws Exception {
        String tableName = "test_count_table_with_wal";
        tEnv.executeSql(
                        String.format(
                                "create table %s ("
                                        + "  id int not null,"
                                        + "  address varchar,"
                                        + "  name varchar,"
                                        + "  primary key (id) NOT ENFORCED)"
                                        + " with ('bucket.num' = '4', 'table.changelog.image' = 'wal')",
                                tableName))
                .await();
        // normal scan
        String query = String.format("SELECT COUNT(*) FROM %s", tableName);
        assertThatThrownBy(() -> tEnv.executeSql(query))
                .hasRootCauseInstanceOf(InvalidTableException.class)
                .hasMessageContaining(
                        "Row count is disabled for this table 'defaultdb.test_count_table_with_wal'.");
    }

    @ParameterizedTest
    @ValueSource(booleans = {true, false})
    void testCountPushDownForLogTable(boolean partitionTable) throws Exception {
        String tableName = partitionTable ? preparePartitionedLogTable() : prepareLogTable();
        int expectedRows = partitionTable ? 10 : 5;
        // normal scan
        String query = String.format("SELECT COUNT(*) FROM %s", tableName);
        assertThat(tEnv.explainSql(query))
                .contains(
                        "aggregates=[grouping=[], aggFunctions=[Count1AggFunction()]]]], fields=[count1$0]");
        CloseableIterator<Row> iterRows = tEnv.executeSql(query).collect();
        List<String> collected = collectRowsWithTimeout(iterRows, 1);
        List<String> expected = Collections.singletonList(String.format("+I[%s]", expectedRows));
        assertThat(collected).isEqualTo(expected);

        // test not push down grouping count.
        assertThatThrownBy(
                        () ->
                                tEnv.explainSql(
                                                String.format(
                                                        "SELECT COUNT(*) FROM %s group by id",
                                                        tableName))
                                        .wait())
                .hasMessageContaining(
                        "Currently, Fluss only support queries on table with datalake enabled or point queries on primary key when it's in batch execution mode.");
    }

    private String prepareSourceTable(String[] keys, String partitionedKey) throws Exception {
        String tableName =
                String.format("test_%s_%s", String.join("_", keys), RandomUtils.nextInt());
        if (partitionedKey == null) {
            tEnv.executeSql(
                    String.format(
                            "create table %s ("
                                    + "  id int not null,"
                                    + "  address varchar,"
                                    + "  name varchar,"
                                    + "  primary key (%s) NOT ENFORCED)"
                                    + " with ('bucket.num' = '4')",
                            tableName, String.join(",", keys)));
        } else {
            tEnv.executeSql(
                    String.format(
                            "create table %s ("
                                    + "  id int not null,"
                                    + "  address varchar,"
                                    + "  name varchar,"
                                    + "  dt varchar,"
                                    + "  primary key (%s) NOT ENFORCED) partitioned by (%s)"
                                    + " with ("
                                    + "  'bucket.num' = '4', "
                                    + "  'table.auto-partition.enabled' = 'true',"
                                    + "  'table.auto-partition.time-unit' = 'year')",
                            tableName, String.join(",", keys), partitionedKey));
        }

        TablePath tablePath = TablePath.of(DEFAULT_DB, tableName);
        String partition1 = null;
        if (partitionedKey != null) {
            Map<Long, String> partitionNameById =
                    waitUntilPartitions(FLUSS_CLUSTER_EXTENSION.getZooKeeperClient(), tablePath);
            // just pick first partition to insert data
            Iterator<String> partitionIterator =
                    partitionNameById.values().stream().sorted().iterator();
            partition1 = partitionIterator.next();
        }

        // prepare table data
        try (Table table = conn.getTable(tablePath)) {
            UpsertWriter upsertWriter = table.newUpsert().createWriter();
            for (int i = 1; i <= 5; i++) {
                Object[] values =
                        partition1 == null
                                ? new Object[] {i, "address" + i, "name" + i}
                                : new Object[] {i, "address" + i, "name" + i, partition1};
                upsertWriter.upsert(row(values));
            }
            upsertWriter.flush();
        }

        return tableName;
    }

    private String prepareLogTable() throws Exception {
        String tableName = String.format("test_log_table_%s", RandomUtils.nextInt());
        tEnv.executeSql(
                String.format(
                        "create table %s ("
                                + "  id int not null,"
                                + "  address varchar,"
                                + "  name varchar)"
                                + " with ("
                                + "  'bucket.num' = '4', "
                                + "  'table.auto-partition.enabled' = 'false' "
                                + ")",
                        tableName));

        TablePath tablePath = TablePath.of(DEFAULT_DB, tableName);

        // prepare table data
        try (Table table = conn.getTable(tablePath)) {
            AppendWriter appendWriter = table.newAppend().createWriter();
            for (int i = 1; i <= 5; i++) {
                Object[] values = new Object[] {i, "address" + i, "name" + i};
                appendWriter.append(row(values));
                // make sure every bucket has records
                appendWriter.flush();
            }
        }

        return tableName;
    }

    protected String preparePartitionedLogTable() throws Exception {
        String tableName = String.format("test_partitioned_log_table_%s", RandomUtils.nextInt());
        tEnv.executeSql(
                String.format(
                        "create table %s ("
                                + "  id int not null,"
                                + "  address varchar,"
                                + "  name varchar,"
                                + "  dt varchar)"
                                + "  partitioned by (dt)"
                                + " with ("
                                + "  'bucket.num' = '4', "
                                + "  'table.auto-partition.enabled' = 'true',"
                                + "  'table.auto-partition.time-unit' = 'year')",
                        tableName));

        TablePath tablePath = TablePath.of(DEFAULT_DB, tableName);
        Map<Long, String> partitionNameById =
                waitUntilPartitions(FLUSS_CLUSTER_EXTENSION.getZooKeeperClient(), tablePath);
        Collection<String> partitions = partitionNameById.values();

        // prepare table data
        try (Table table = conn.getTable(tablePath)) {
            AppendWriter appendWriter = table.newAppend().createWriter();
            for (int i = 1; i <= 5; i++) {
                for (String partition : partitions) {
                    Object[] values = new Object[] {i, "address" + i, "name" + i, partition};
                    appendWriter.append(row(values));
                    // make sure every bucket has records
                    appendWriter.flush();
                }
            }
        }

        return tableName;
    }
}
