package com.lanyuanxiaoyao.service.executor.task;

import cn.hutool.core.collection.IterUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import de.vandermeer.asciitable.AsciiTable;
import de.vandermeer.asciitable.CWC_LongestLine;
import de.vandermeer.asciithemes.TA_GridThemes;
import java.util.function.Function;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.EnvironmentSettings;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.TableEnvironment;
import org.apache.flink.table.api.TableResult;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.types.DataType;
import org.apache.flink.types.Row;
import org.apache.flink.util.CloseableIterator;
import org.apache.hudi.common.table.HoodieTableConfig;
import org.apache.hudi.common.table.HoodieTableMetaClient;
import org.apache.hudi.common.table.TableSchemaResolver;
import org.apache.hudi.org.apache.avro.Schema;
import org.apache.hudi.util.AvroSchemaConverter;
import org.apache.hudi.util.HoodiePipeline;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.factory.Maps;
import org.eclipse.collections.api.list.ImmutableList;
import org.eclipse.collections.api.map.ImmutableMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 通过执行Flink SQL查数据
 *
 * @author lanyuanxiaoyao
 * @date 2024-05-17
 */
public class SQLExecutor {
    private static final Logger logger = LoggerFactory.getLogger(SQLExecutor.class);

    protected static CloseableIterator<Row> executeStream(HoodieTableMetaClient metaClient, Function<Table, TableResult> SQL) throws Exception {
        HoodieTableConfig tableConfig = metaClient.getTableConfig();
        Schema schema = new TableSchemaResolver(metaClient).getTableAvroSchema();

        StreamExecutionEnvironment executionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment();
        StreamTableEnvironment tableEnvironment = StreamTableEnvironment.create(executionEnvironment);
        HoodiePipeline.Builder builder = HoodiePipeline.builder(tableConfig.getTableName());
        for (Schema.Field field : schema.getFields()) {
            DataType type = AvroSchemaConverter.convertToDataType(field.schema());
            builder.column(StrUtil.format("{} {}", field.name(), type.getLogicalType()));
        }
        tableConfig.getPartitionFields().ifPresent(builder::partition);
        tableConfig.getRecordKeyFields().ifPresent(builder::pk);
        builder.option("connector", "hudi");
        builder.options(IterUtil.toMap(metaClient.getHadoopConf()));

        Table table = tableEnvironment.fromDataStream(builder.source(executionEnvironment));
        return SQL.apply(table).collect();
    }

    protected static CloseableIterator<Row> executeBatch(HoodieTableMetaClient metaClient, Function<String, String> SQL) throws Exception {
        HoodieTableConfig tableConfig = metaClient.getTableConfig();
        Schema schema = new TableSchemaResolver(metaClient).getTableAvroSchema();

        TableEnvironment tableEnvironment = TableEnvironment.create(
                EnvironmentSettings.newInstance()
                        .inBatchMode()
                        .useBlinkPlanner()
                        .build()
        );
        tableEnvironment.executeSql(createHoodieTableDDL(
                tableConfig.getTableName(),
                Lists.immutable.ofAll(schema.getFields())
                        .collect(field -> {
                            DataType type = AvroSchemaConverter.convertToDataType(field.schema());
                            return StrUtil.format("{} {}", field.name(), type.getLogicalType());
                        }),
                tableConfig.getPartitionFields().map(Lists.immutable::of).orElse(Lists.immutable.empty()),
                tableConfig.getRecordKeyFields().map(Lists.immutable::of).orElse(Lists.immutable.empty()),
                Maps.immutable.ofAll(IterUtil.toMap(metaClient.getHadoopConf()))
        ));
        String sql = SQL.apply(tableConfig.getTableName());
        logger.info("SQL: {}", sql);
        return tableEnvironment.executeSql(sql).collect();
    }

    protected static String generateResult(CloseableIterator<Row> iterator, ImmutableList<String> fields) {
        AsciiTable table = new AsciiTable();
        table.addRule();
        table.addRow(fields.toArray());
        table.addRule();
        iterator.forEachRemaining(row -> {
            table.addRow(fields.collect(name -> {
                Object field = row.getField(name);
                return ObjectUtil.isNull(field) ? "null" : field;
            }).toArray());
        });
        table.addRule();
        table.getRenderer().setCWC(new CWC_LongestLine());
        table.getContext().setGridTheme(TA_GridThemes.NONE);
        table.setPaddingLeftRight(1);
        return table.render();
    }

    private static String createHoodieTableDDL(
            String tableName,
            ImmutableList<String> fields,
            ImmutableList<String> primaryKeys,
            ImmutableList<String> partitionKeys,
            ImmutableMap<String, String> options
    ) {
        StringBuilder builder = new StringBuilder();
        builder.append("create table if not exists ")
                .append(tableName)
                .append("(\n");
        for (String field : fields) {
            builder.append("  ")
                    .append(field)
                    .append(",\n");
        }
        builder.append("  PRIMARY KEY(")
                .append(primaryKeys.makeString(","))
                .append(") NOT ENFORCED\n")
                .append(")\n");
        if (!partitionKeys.isEmpty()) {
            String partitons = partitionKeys
                    .collect(key -> StrUtil.format("`{}`", key))
                    .makeString(",");
            builder.append("PARTITIONED BY (")
                    .append(partitons)
                    .append(")\n");
        }
        builder.append("with ('connector' = 'hudi'");
        options.forEachKeyValue((k, v) -> builder
                .append(",\n")
                .append("  '")
                .append(k)
                .append("' = '")
                .append(v)
                .append("'"));
        builder.append("\n)");
        return builder.toString();
    }
}
