package com.easyagents.federation.sql.adapter.jdbc;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.execute.FederationFragmentExplainContext;
import com.easyagents.federation.sql.execute.FederationFragmentExplainer;
import com.easyagents.federation.sql.execute.FederationPhysicalExplain;
import com.easyagents.federation.sql.execute.SqlParameter;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.SQLTimeoutException;
import java.sql.SQLTransientConnectionException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;

/**
 * MySQL、PostgreSQL 和 H2 的非 ANALYZE 物理 Explain 实现。
 */
final class JdbcFederationFragmentExplainer implements FederationFragmentExplainer {

    private static final ObjectMapper JSON = new ObjectMapper();

    /** {@inheritDoc} */
    @Override
    public FederationPhysicalExplain explain(FederationFragmentExplainContext context) {
        String product = normalize(context.compatibility().databaseProduct());
        String explainSql = explainSql(product, context.sql());
        if (explainSql == null) {
            return FederationPhysicalExplain.unavailable(
                "physical Explain is not implemented for "
                    + context.compatibility().databaseProduct()
            );
        }
        Connection acquired = acquireConnection(context);
        try (Connection connection = acquired) {
            context.executionGuard().ensureAllowed();
            if (!connection.isReadOnly()) {
                connection.setReadOnly(true);
            }
            try (PreparedStatement statement = connection.prepareStatement(explainSql)) {
                int queryTimeout = context.executionGuard().boundedQueryTimeoutSeconds(
                    context.queryTimeoutSeconds()
                );
                if (queryTimeout > 0) {
                    statement.setQueryTimeout(queryTimeout);
                }
                bind(statement, context.parameters());
                context.executionGuard().ensureAllowed();
                try (ResultSet resultSet = statement.executeQuery()) {
                    context.executionGuard().ensureAllowed();
                    String nativePlan = readPlan(resultSet);
                    return normalizePlan(product, nativePlan);
                }
            }
        } catch (SQLException | RuntimeException exception) {
            if (exception instanceof FederationSqlException federationSqlException) {
                throw federationSqlException;
            }
            throw new FederationSqlException(
                FederationSqlErrorCode.EXPLAIN_FAILED,
                "physical database Explain failed",
                exception
            );
        }
    }

    /**
     * 在统一截止时间约束下获取物理 Explain 连接。
     *
     * @param context 分片 Explain 上下文
     * @return 已获取连接
     * @throws FederationSqlException 获取超时、失败或查询已终止时抛出
     */
    private static Connection acquireConnection(FederationFragmentExplainContext context) {
        try {
            context.executionGuard().ensureAllowed();
            return context.dataSource().getConnection();
        } catch (SQLException | RuntimeException exception) {
            if (exception instanceof FederationSqlException federationSqlException) {
                throw federationSqlException;
            }
            context.executionGuard().ensureAllowed();
            FederationSqlErrorCode code = exception instanceof SQLTimeoutException
                || exception instanceof SQLTransientConnectionException
                ? FederationSqlErrorCode.CONNECTION_ACQUISITION_TIMEOUT
                : FederationSqlErrorCode.CONNECTION_ACQUISITION_FAILED;
            String message = code == FederationSqlErrorCode.CONNECTION_ACQUISITION_TIMEOUT
                ? "timed out while acquiring a JDBC connection for physical Explain"
                : "failed to acquire a JDBC connection for physical Explain";
            throw new FederationSqlException(code, message, exception);
        }
    }

    private static String explainSql(String product, String sql) {
        if (product.contains("mysql")) {
            return "EXPLAIN FORMAT=JSON " + sql;
        }
        if (product.contains("postgres")) {
            return "EXPLAIN (FORMAT JSON, ANALYZE FALSE, COSTS TRUE, VERBOSE FALSE, BUFFERS FALSE) "
                + sql;
        }
        if (product.equals("h2")) {
            return "EXPLAIN " + sql;
        }
        return null;
    }

    private static void bind(PreparedStatement statement, List<SqlParameter> parameters)
        throws SQLException {
        for (int index = 0; index < parameters.size(); index++) {
            SqlParameter parameter = parameters.get(index);
            if (parameter.value() == null) {
                statement.setNull(index + 1, parameter.jdbcType());
            } else {
                statement.setObject(index + 1, parameter.value(), parameter.jdbcType());
            }
        }
    }

    private static String readPlan(ResultSet resultSet) throws SQLException {
        StringBuilder plan = new StringBuilder();
        ResultSetMetaData metadata = resultSet.getMetaData();
        while (resultSet.next()) {
            if (!plan.isEmpty()) {
                plan.append('\n');
            }
            for (int column = 1; column <= metadata.getColumnCount(); column++) {
                if (column > 1) {
                    plan.append('\t');
                }
                Object value = resultSet.getObject(column);
                if (value != null) {
                    plan.append(value);
                }
            }
        }
        return plan.toString();
    }

    private static FederationPhysicalExplain normalizePlan(String product, String nativePlan) {
        if (!nativePlan.isBlank() && (product.contains("mysql") || product.contains("postgres"))) {
            try {
                JsonNode root = JSON.readTree(nativePlan);
                return product.contains("mysql")
                    ? normalizeMysql(root, nativePlan)
                    : normalizePostgresql(root, nativePlan);
            } catch (Exception ignored) {
                // 原生计划仍可用；归一化失败不会伪造索引结论。
            }
        }
        return new FederationPhysicalExplain(
            true,
            nativePlan,
            null,
            null,
            List.of(),
            null,
            null,
            null,
            "native plan is available; normalized index fields are unavailable"
        );
    }

    private static FederationPhysicalExplain normalizeMysql(JsonNode root, String nativePlan) {
        JsonNode table = findObjectWithField(root, "access_type");
        if (table == null) {
            return nativeOnly(nativePlan, "MySQL plan contains no normalized table access node");
        }
        List<String> candidates = stringValues(table.get("possible_keys"));
        return new FederationPhysicalExplain(
            true,
            nativePlan,
            "table",
            text(table, "access_type"),
            candidates,
            text(table, "key"),
            longValue(table, "rows_examined_per_scan", "rows"),
            firstText(table, "attached_condition", "index_condition"),
            "normalized from MySQL JSON Explain"
        );
    }

    private static FederationPhysicalExplain normalizePostgresql(JsonNode root, String nativePlan) {
        JsonNode plan = root.isArray() && !root.isEmpty() ? root.get(0).get("Plan") : root.get("Plan");
        JsonNode scan = findObjectWithField(plan, "Index Name");
        if (scan == null) {
            scan = findObjectWithTextSuffix(plan, "Node Type", "Scan");
        }
        if (scan == null) {
            return nativeOnly(nativePlan, "PostgreSQL plan contains no normalized plan node");
        }
        return new FederationPhysicalExplain(
            true,
            nativePlan,
            text(scan, "Node Type"),
            text(scan, "Node Type"),
            List.of(),
            text(scan, "Index Name"),
            longValue(scan, "Plan Rows"),
            firstText(scan, "Index Cond", "Filter", "Join Filter"),
            "normalized from PostgreSQL JSON Explain"
        );
    }

    private static FederationPhysicalExplain nativeOnly(String nativePlan, String diagnostic) {
        return new FederationPhysicalExplain(
            true,
            nativePlan,
            null,
            null,
            List.of(),
            null,
            null,
            null,
            diagnostic
        );
    }

    private static JsonNode findObjectWithField(JsonNode node, String field) {
        if (node == null) {
            return null;
        }
        if (node.isObject() && node.has(field)) {
            return node;
        }
        Iterator<JsonNode> children = node.elements();
        while (children.hasNext()) {
            JsonNode found = findObjectWithField(children.next(), field);
            if (found != null) {
                return found;
            }
        }
        return null;
    }

    private static JsonNode findObjectWithTextSuffix(
        JsonNode node,
        String field,
        String suffix
    ) {
        if (node == null) {
            return null;
        }
        if (node.isObject()) {
            String value = text(node, field);
            if (value != null && value.endsWith(suffix)) {
                return node;
            }
        }
        Iterator<JsonNode> children = node.elements();
        while (children.hasNext()) {
            JsonNode found = findObjectWithTextSuffix(children.next(), field, suffix);
            if (found != null) {
                return found;
            }
        }
        return null;
    }

    private static List<String> stringValues(JsonNode node) {
        if (node == null || node.isNull()) {
            return List.of();
        }
        if (node.isArray()) {
            List<String> values = new ArrayList<>();
            node.forEach(value -> values.add(value.asText()));
            return List.copyOf(values);
        }
        return List.of(node.asText());
    }

    private static String firstText(JsonNode node, String... fields) {
        for (String field : fields) {
            String value = text(node, field);
            if (value != null) {
                return value;
            }
        }
        return null;
    }

    private static String text(JsonNode node, String field) {
        JsonNode value = node == null ? null : node.get(field);
        return value == null || value.isNull() ? null : value.asText();
    }

    private static Long longValue(JsonNode node, String... fields) {
        for (String field : fields) {
            JsonNode value = node == null ? null : node.get(field);
            if (value != null && value.isNumber()) {
                return value.longValue();
            }
        }
        return null;
    }

    private static String normalize(String product) {
        return product == null ? "" : product.trim().toLowerCase(Locale.ROOT);
    }
}
