package com.purui.prims.headless.core.translator.calcite.sql.render;

import com.purui.prims.common.pojo.JoinCondition;
import com.purui.prims.headless.api.pojo.ModelDetail;
import com.purui.prims.headless.api.pojo.TableJoinRelation;
import com.purui.prims.headless.api.pojo.enums.EngineType;
import com.purui.prims.headless.api.pojo.response.DatabaseResp;
import com.purui.prims.headless.core.translator.calcite.Configuration;
import com.purui.prims.headless.core.translator.calcite.schema.SemanticSchema;
import com.purui.prims.headless.core.translator.calcite.sql.S2SQLSqlValidatorImpl;
import com.purui.prims.headless.core.translator.calcite.sql.node.SemanticNode;
import lombok.extern.slf4j.Slf4j;
import org.apache.calcite.jdbc.CalciteSchema;
import org.apache.calcite.prepare.CalciteCatalogReader;
import org.apache.calcite.prepare.Prepare;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.sql.JoinConditionType;
import org.apache.calcite.sql.JoinType;
import org.apache.calcite.sql.SqlBasicCall;
import org.apache.calcite.sql.SqlIdentifier;
import org.apache.calcite.sql.SqlJoin;
import org.apache.calcite.sql.SqlLiteral;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlNodeList;
import org.apache.calcite.sql.SqlSelect;
import org.apache.calcite.sql.SqlTableRef;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.parser.SqlParserPos;
import org.apache.calcite.sql.validate.ParameterScope;
import org.apache.calcite.sql.validate.SqlValidatorScope;
import org.apache.commons.lang3.StringUtils;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 动态Join类
 *
 * @author charlie
 * @date 2024/7/4
 */
@Slf4j
public class ModelSQLRender {
    /**
     * 根据模型设置，生成动态SQL语句
     *
     * @param modelDetail
     * @param database
     * @return
     * @throws Exception
     */
    public String createJoinQuery(ModelDetail modelDetail, DatabaseResp database) throws Exception {
        if (StringUtils.isNotEmpty(modelDetail.getSqlQuery()) || CollectionUtils.isEmpty(modelDetail.getJoinRelations())) {
            return modelDetail.getSqlQuery();
        }
        String leftTable = modelDetail.getTableQuery();
        EngineType engineType = EngineType.fromString(database.getType());
        SemanticSchema schema = SemanticSchema.newBuilder(modelDetail.getTableQuery()).build();
        SqlValidatorScope scope = getScope(schema, engineType);

        return createDynamicLeftJoinQuery(modelDetail.getJoinRelations(), leftTable, scope, engineType);
    }

    // 默认采用LeftJoin
    public String createDynamicLeftJoinQuery(List<TableJoinRelation> tableJoins, String primaryTableName, SqlValidatorScope scope,
                                             EngineType engineType)
        throws Exception {
        String primaryAlias = "m";
        SqlNode mainTableRef = getTableReference(primaryTableName, primaryAlias);
        Map<String, String> aliasMap = new HashMap<>();
        aliasMap.put(primaryTableName, primaryAlias);
        int i = 1;
        boolean exists = tableJoins.stream().anyMatch(join -> primaryTableName.equals(join.getLeft()));
        if (exists) {
            mainTableRef = createAllJoin(mainTableRef, aliasMap, tableJoins, primaryTableName, i, scope, engineType);

        }

        // 构建SELECT语句，选择所有字段并加上别名前缀
        SqlNodeList selectList = new SqlNodeList(SqlParserPos.ZERO);
        for (TableJoinRelation tj : tableJoins) {
            String alias = aliasMap.get(tj.getLeft());
            for (String field : tj.getLeftFields()) {
                selectList.add(new SqlIdentifier(alias + "." + field, SqlParserPos.ZERO));
            }
            alias = aliasMap.get(tj.getRight());
            for (String field : tj.getRightFields()) {
                selectList.add(new SqlIdentifier(alias + "." + field, SqlParserPos.ZERO));
            }
        }

        SqlSelect select = new SqlSelect(SqlParserPos.ZERO, null, selectList, mainTableRef, null, null, null, null, null, null, null, null);

        // 格式化输出SQL
        return SemanticNode.getSql(select, engineType);
    }

    // 递归处理JOIN关系
    private SqlNode createAllJoin(SqlNode mainTableRef, Map<String, String> aliasMap, List<TableJoinRelation> tableJoins,
                                  String primaryTableName, int i, SqlValidatorScope scope, EngineType engineType) throws Exception {
        String finalPrimaryTableName = primaryTableName;
        List<TableJoinRelation> joins =
            tableJoins.stream().filter(j -> finalPrimaryTableName.equals(j.getLeft())).collect(Collectors.toList());
        for (TableJoinRelation join : joins) {
            mainTableRef = createOneJoin(mainTableRef, aliasMap, join, primaryTableName, scope, engineType, i);
            i++;
        }
        // 遍历所有的JOIN表
        for (TableJoinRelation join : joins) {
            String finalPrimaryTableName1 = join.getRight();
            List<TableJoinRelation> nextJoins =
                tableJoins.stream().filter(j -> finalPrimaryTableName1.equals(j.getLeft())).collect(Collectors.toList());
            if (!CollectionUtils.isEmpty(nextJoins)) {
                primaryTableName = join.getRight();
                return createAllJoin(mainTableRef, aliasMap, nextJoins, primaryTableName, i, scope, engineType);
            }
        }
        return mainTableRef;
    }

    private SqlNode createOneJoin(SqlNode mainTableRef, Map<String, String> aliasMap, TableJoinRelation join, String primaryTableName,
                                  SqlValidatorScope scope, EngineType engineType, int i)
        throws Exception {
        SqlNode joinTableRef = getTableReference(join.getRight(), "j" + i);
        aliasMap.put(join.getRight(), "j" + i);
        SqlNode joinCondition =
            getCondition(aliasMap.get(join.getLeft()), aliasMap.get(join.getRight()), join.getConditions(), scope, engineType);
        mainTableRef = new SqlJoin(SqlParserPos.ZERO,
            mainTableRef,
            SqlLiteral.createBoolean(false, SqlParserPos.ZERO),
            JoinType.LEFT.symbol(SqlParserPos.ZERO),
            joinTableRef,
            JoinConditionType.ON.symbol(SqlParserPos.ZERO),
            joinCondition);

        return mainTableRef;
    }

    private SqlNode getTableReference(String tableName, String alias) throws Exception {
        SqlNode table = new SqlTableRef(SqlParserPos.ZERO,
            new SqlIdentifier(tableName, SqlParserPos.ZERO), null);
        return SemanticNode.buildAs(alias, table);
    }

    private SqlValidatorScope getScope(SemanticSchema schema, EngineType engineType) {
        Map<String, RelDataType> nameToTypeMap = new HashMap<>();
        CalciteSchema rootSchema = CalciteSchema.createRootSchema(true, false);
        rootSchema.add(schema.getSchemaKey(), schema);
        Prepare.CatalogReader catalogReader = new CalciteCatalogReader(
            rootSchema,
            Collections.singletonList(schema.getSchemaKey()),
            Configuration.typeFactory,
            Configuration.config
        );
        S2SQLSqlValidatorImpl s2SQLSqlValidator = new S2SQLSqlValidatorImpl(Configuration.operatorTable, catalogReader,
            Configuration.typeFactory, Configuration.getValidatorConfig(engineType));
        return new ParameterScope(s2SQLSqlValidator, nameToTypeMap);
    }

    private SqlNode getCondition(String left, String right, List<JoinCondition> conditions, SqlValidatorScope scope, EngineType engineType)
        throws Exception {
        SqlNode condition = null;
        for (JoinCondition on : conditions) {
            List<SqlNode> ons = new ArrayList<>();
            ons.add(SemanticNode.parse(left + "." + on.getLeftField(), scope, engineType));
            ons.add(SemanticNode.parse(right + "." + on.getRightField(), scope, engineType));
            if (condition == null) {
                condition = new SqlBasicCall(
                    SqlStdOperatorTable.EQUALS,
                    ons,
                    SqlParserPos.ZERO, null);
                continue;
            }
            SqlNode addCondition = new SqlBasicCall(
                SqlStdOperatorTable.EQUALS,
                ons,
                SqlParserPos.ZERO, null);
            condition = new SqlBasicCall(
                SqlStdOperatorTable.AND,
                new ArrayList<>(Arrays.asList(condition, addCondition)),
                SqlParserPos.ZERO, null);
        }
        return condition;
    }
}
