package com.agent.management.rag.structured;

import com.agent.management.rag.api.RagRetriever;
import com.agent.management.rag.model.RagEvidence;
import com.agent.management.rag.model.RagQuery;
import com.agent.management.rag.model.RagRetrievalResult;
import com.agent.management.rag.model.RagSourceRef;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.service.DataSourceService;
import com.agent.management.service.SchemaExplorerService;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.regex.Pattern;

@Service
@RequiredArgsConstructor
public class StructuredDataRagRetriever implements RagRetriever {
    private final ExplicitSqlGenerationService explicit;
    private final VannaSqlGenerationService generated;
    private final StructuredSqlRelationshipValidator relationshipValidator;
    private final StructuredQueryContextService contexts;
    private final SchemaExplorerService explorer;
    private final DataSourceService dataSources;

    public RagSourceType sourceType() { return RagSourceType.STRUCTURED_DATA; }

    public RagRetrievalResult retrieve(RagQuery query) {
        RagRetrievalResult out = new RagRetrievalResult();
        out.setSourceType(sourceType());
        try {
            Long id = sourceId(query);
            Optional<String> explicitSql = explicit.generateSql(query, id);
            boolean auto = explicitSql.isEmpty();
            String sql;
            StructuredQueryContext context = null;
            List<String> warnings = List.of();
            List<String> usedRelationships = List.of();
            List<String> relationshipWarnings = List.of();
            int repairCount = 0;

            if (auto) {
                if (!isMysqlSource(id)) {
                    addStructuredDiagnostic(out, id, "UNSUPPORTED_SQL_DIALECT",
                            "自动 Text2SQL 当前仅支持 MySQL；该数据源方言暂不支持自动生成 SQL。", Map.of());
                    return out;
                }
                try {
                    context = contexts.build(query, id);
                } catch (Exception error) {
                    out.getDiagnostics().put("error", "structured query context failed: " + error.getMessage());
                    return out;
                }
                Optional<MetricDiagnostic> metricDiagnostic = metricDiagnostic(query, context);
                if (metricDiagnostic.isPresent()) {
                    addStructuredDiagnostic(out, id, metricDiagnostic.get().code(), metricDiagnostic.get().message(),
                            metricDiagnostic.get().metadata());
                    return out;
                }
                VannaSqlGenerationService.GeneratedSql generation;
                try {
                    Optional<VannaSqlGenerationService.GeneratedSql> generatedSql = generated.generate(query, id, context);
                    if (generatedSql.isEmpty()) {
                        out.getDiagnostics().put("warning", "no SQL provided and automatic generation is disabled or returned empty");
                        return out;
                    }
                    generation = generatedSql.get();
                } catch (Exception error) {
                    out.getDiagnostics().put("error", "SQL generation failed: " + error.getMessage());
                    return out;
                }
                sql = generation.sql();
                context = generation.context();
                if (!id.equals(context.dataSourceId())) {
                    throw new IllegalStateException("generated SQL context belongs to another datasource");
                }
                warnings = generation.warnings();
                PreflightResult preflight = preflightAndRepair(query, id, sql, context, out);
                if (preflight.sql().isEmpty()) return out;
                sql = preflight.sql().get();
                repairCount = preflight.repairCount();
                usedRelationships = preflight.usedRelationships();
                relationshipWarnings = preflight.relationshipWarnings();
            } else {
                sql = explicitSql.get();
                try {
                    explorer.explainQuery(id, sql);
                } catch (Exception error) {
                    out.getDiagnostics().put("error", "explicit SQL EXPLAIN failed: " + error.getMessage());
                    return out;
                }
            }

            int limit = maxRows(query);
            SchemaExplorerService.QueryResult result = explorer.executeQuery(id, sql, limit);
            var source = dataSources.get(id);
            RagEvidence evidence = new RagEvidence();
            evidence.setId("sql-" + UUID.randomUUID());
            evidence.setSourceType(sourceType());
            evidence.setEvidenceType("SQL_RESULT");
            evidence.setTitle(source.getName());
            evidence.setSourceName(evidence.getTitle());
            evidence.setContent(summary(result));
            evidence.setScore(1.0);
            RagSourceRef ref = new RagSourceRef();
            ref.setSourceId(String.valueOf(id));
            ref.setSourceName(evidence.getTitle());
            ref.setLocator(Map.of("datasourceId", id));
            evidence.setSourceRef(ref);
            evidence.setPayload(Map.of("columns", result.columns(), "rows", result.rows(), "rowCount", result.rows().size()));
            Map<String, Object> metadata = new LinkedHashMap<>();
            metadata.put("sql", sql);
            metadata.put(auto ? "generatedSql" : "explicitSql", sql);
            metadata.put("executionTimeMs", result.executionTimeMs());
            metadata.put("dataSourceId", id);
            metadata.put("warnings", warnings);
            metadata.put("dialect", context == null
                    ? (source.getType() == null ? "" : source.getType().toLowerCase(Locale.ROOT))
                    : context.dialect());
            metadata.put("schemaVersion", context == null ? "" : context.schemaVersion());
            metadata.put("schemaFingerprint", context == null ? "" : context.schemaFingerprint());
            metadata.put("profileVersion", context == null ? "" : context.profileVersion());
            metadata.put("selectedTables", context == null ? List.of() : context.selectedTables());
            metadata.put("sampledValueFields", context == null ? List.of() : context.sampledValueFields());
            metadata.put("referenceTables", context == null ? List.of() : context.referenceTables());
            metadata.put("contextWarnings", context == null ? List.of() : context.warnings());
            metadata.put("usedRelationships", usedRelationships);
            metadata.put("relationshipWarnings", relationshipWarnings);
            metadata.put("repaired", repairCount > 0);
            metadata.put("repairCount", repairCount);
            evidence.setMetadata(metadata);
            out.getEvidences().add(evidence);
        } catch (Exception error) {
            out.getDiagnostics().put("error", error.getMessage());
        }
        return out;
    }

    private PreflightResult preflightAndRepair(RagQuery query, Long sourceId, String sql,
                                               StructuredQueryContext context, RagRetrievalResult out) {
        try {
            StructuredSqlRelationshipValidator.ValidationResult validation = relationshipValidator.validate(sql, context);
            explorer.explainQuery(sourceId, sql);
            return new PreflightResult(Optional.of(sql), 0,
                    validation.usedRelationships(), validation.warnings());
        } catch (Exception preflightError) {
            try {
                Optional<String> repaired = generated.repair(query, sourceId, sql, preflightError.getMessage(), context);
                if (repaired.isEmpty()) throw new IllegalArgumentException("repair returned empty SQL");
                StructuredSqlRelationshipValidator.ValidationResult validation =
                        relationshipValidator.validate(repaired.get(), context);
                explorer.explainQuery(sourceId, repaired.get());
                out.getDiagnostics().put("repair", "generated SQL failed EXPLAIN and was repaired once");
                return new PreflightResult(repaired, 1,
                        validation.usedRelationships(), validation.warnings());
            } catch (Exception repairError) {
                out.getDiagnostics().put("error", "SQL preflight/repair failed: " + repairError.getMessage());
                return new PreflightResult(Optional.empty(), 1, List.of(), List.of(preflightError.getMessage()));
            }
        }
    }

    private static int maxRows(RagQuery query) {
        Object configured = query.getFilters() == null ? null : query.getFilters().get("maxRows");
        int value = configured instanceof Number number ? number.intValue()
                : (query.getTopK() != null ? query.getTopK() : 50);
        return Math.max(1, Math.min(value, 1000));
    }

    private static Long sourceId(RagQuery query) {
        if (query.getSourceIds() == null || query.getSourceIds().isEmpty()) {
            throw new IllegalArgumentException("sourceIds must contain datasourceId");
        }
        if (query.getSourceIds().size() != 1) {
            throw new IllegalArgumentException("structured Text2SQL supports exactly one datasourceId per request");
        }
        return Long.valueOf(query.getSourceIds().get(0));
    }

    private void addStructuredDiagnostic(RagRetrievalResult out, Long sourceId, String code,
                                         String message, Map<String, Object> metadata) {
        out.getDiagnostics().put("warning", message);
        var source = dataSources.get(sourceId);
        RagEvidence evidence = new RagEvidence();
        evidence.setId("sql-diagnostic-" + UUID.randomUUID());
        evidence.setSourceType(sourceType());
        evidence.setEvidenceType("SQL_DIAGNOSTIC");
        evidence.setTitle(source.getName());
        evidence.setSourceName(evidence.getTitle());
        evidence.setContent(message);
        evidence.setScore(1.0);
        RagSourceRef ref = new RagSourceRef();
        ref.setSourceId(String.valueOf(sourceId));
        ref.setSourceName(evidence.getTitle());
        ref.setLocator(Map.of("datasourceId", sourceId));
        evidence.setSourceRef(ref);
        evidence.setPayload(Map.of("columns", List.of(), "rows", List.of(), "rowCount", 0));
        Map<String, Object> result = new LinkedHashMap<>(metadata == null ? Map.of() : metadata);
        result.put("diagnostic", code);
        result.put("dataSourceId", sourceId);
        evidence.setMetadata(result);
        out.getEvidences().add(evidence);
    }

    private boolean isMysqlSource(Long sourceId) {
        var source = dataSources.get(sourceId);
        return source.getType() != null && "MYSQL".equalsIgnoreCase(source.getType());
    }

    private Optional<MetricDiagnostic> metricDiagnostic(RagQuery query, StructuredQueryContext context) {
        if (!isMetricQuestion(query.getQuery())) return Optional.empty();
        List<String> metricFields = metricFields(context);
        if (metricFields.isEmpty()) {
            return Optional.of(new MetricDiagnostic("NO_STRUCTURED_METRIC_DATA",
                    "结构化库无可支撑该指标的预计算数据；当前画像未发现评分、排名或优先级字段。",
                    Map.of("missing", "metric_fields", "selectedTables", context.selectedTables())));
        }
        List<String> mentions = entityMentions(query, context);
        if (!mentions.isEmpty()
                && !mentionsGroundedForMetricQuery(mentions, metricFields, context)
                && !mentionsGroundedInScopedEntityFields(mentions, context)
                && !mentionsExistInScopedEntityFields(mentions, context)) {
            return Optional.of(new MetricDiagnostic("NO_STRUCTURED_METRIC_DATA",
                    "结构化库无可支撑该指标的预计算数据；问题中的实体或候选值未出现在含指标字段的数据表中。",
                    Map.of("missing", "metric_entity_values", "entityMentions", mentions,
                            "metricFields", metricFields, "selectedTables", context.selectedTables())));
        }
        return Optional.empty();
    }

    private static boolean isMetricQuestion(String question) {
        String text = question == null ? "" : question.toLowerCase(Locale.ROOT);
        return containsAny(text, "评分", "得分", "score", "final_score", "排名", "排序", "优先级",
                "优先", "综合", "评估", "top", "best", "highest", "rank", "rating");
    }

    private static List<String> metricFields(StructuredQueryContext context) {
        List<String> result = new java.util.ArrayList<>();
        for (String rawLine : context.minimalDdl().split("\\R")) {
            String line = rawLine.trim();
            java.util.regex.Matcher table = Pattern.compile("(?i)^TABLE\\s+`?([A-Za-z_][A-Za-z0-9_]*)`?\\s*\\((.*)\\)").matcher(line);
            if (!table.find()) continue;
            String tableName = table.group(1);
            for (String rawColumn : table.group(2).split(",")) {
                java.util.regex.Matcher column = Pattern.compile("\\s*`?([A-Za-z_][A-Za-z0-9_]*)`?\\s+").matcher(rawColumn);
                if (!column.find()) continue;
                String name = column.group(1).toLowerCase(Locale.ROOT);
                if (name.equals("score") || name.endsWith("_score") || name.contains("rank")
                        || name.contains("priority") || name.contains("rating")
                        || name.equals("weight") || name.equals("level")) {
                    result.add(tableName + "." + column.group(1));
                }
            }
        }
        return result;
    }

    private static List<String> entityMentions(RagQuery query, StructuredQueryContext context) {
        LinkedHashSet<String> result = new LinkedHashSet<>();
        if (query.getQuery() != null) {
            java.util.regex.Matcher matcher = Pattern.compile("(?:[A-Z][A-Z0-9]*_)+(?:[A-Z0-9]+)").matcher(query.getQuery());
            while (matcher.find()) result.add(matcher.group());
        }
        return result.stream().toList();
    }

    private static boolean mentionsGroundedForMetricQuery(List<String> mentions, List<String> metricFields,
                                                          StructuredQueryContext context) {
        Set<String> metricTables = metricFields.stream()
                .map(field -> field.substring(0, field.indexOf('.')))
                .collect(java.util.stream.Collectors.toCollection(LinkedHashSet::new));
        String evidence = metricReachableTables(metricTables, context).stream()
                .map(table -> valuesForTable(table, context))
                .reduce("", (left, right) -> left + " " + right)
                .toLowerCase(Locale.ROOT);
        return mentions.stream().map(value -> value.toLowerCase(Locale.ROOT)).anyMatch(evidence::contains);
    }

    private static Set<String> metricReachableTables(Set<String> metricTables, StructuredQueryContext context) {
        Set<String> result = new LinkedHashSet<>(metricTables);
        for (StructuredRelationship relationship : context.relationships()) {
            if (metricTables.contains(relationship.sourceTable())) result.add(relationship.targetTable());
            if (metricTables.contains(relationship.targetTable())) result.add(relationship.sourceTable());
        }
        for (var foreignKey : context.foreignKeys()) {
            if (metricTables.contains(foreignKey.sourceTable())) result.add(foreignKey.targetTable());
            if (metricTables.contains(foreignKey.targetTable())) result.add(foreignKey.sourceTable());
        }
        return result;
    }

    private static boolean mentionsGroundedInScopedEntityFields(List<String> mentions, StructuredQueryContext context) {
        StringBuilder evidence = new StringBuilder();
        context.sampledValues().forEach((field, values) -> {
            if (isScopedEntityField(field)) values.forEach(value -> evidence.append(' ').append(value));
        });
        String text = evidence.toString().toLowerCase(Locale.ROOT);
        return mentions.stream().map(value -> value.toLowerCase(Locale.ROOT)).anyMatch(text::contains);
    }

    private static boolean isScopedEntityField(String field) {
        String value = field == null ? "" : field.toLowerCase(Locale.ROOT);
        int index = value.indexOf('.');
        String table = index < 0 ? "" : value.substring(0, index);
        String column = index < 0 ? value : value.substring(index + 1);
        return table.contains("scenario") || table.contains("mission")
                || column.contains("target") || column.contains("candidate")
                || column.contains("platform") || column.contains("scenario")
                || column.contains("mission");
    }

    private boolean mentionsExistInScopedEntityFields(List<String> mentions, StructuredQueryContext context) {
        for (ScopedField field : scopedEntityFields(context)) {
            for (String mention : mentions) {
                String sql = "SELECT 1 FROM " + quoteIdentifier(field.table()) + " WHERE "
                        + quoteIdentifier(field.column()) + " = " + sqlLiteral(mention) + " LIMIT 1";
                try {
                    SchemaExplorerService.QueryResult result = explorer.executeQuery(context.dataSourceId(), sql, 1);
                    if (!result.rows().isEmpty()) return true;
                } catch (Exception ignored) {
                    // A failed grounding probe should not break retrieval; the normal diagnostic will explain the miss.
                }
            }
        }
        return false;
    }

    private static List<ScopedField> scopedEntityFields(StructuredQueryContext context) {
        List<ScopedField> result = new java.util.ArrayList<>();
        for (String rawLine : context.minimalDdl().split("\\R")) {
            String line = rawLine.trim();
            java.util.regex.Matcher table = Pattern.compile("(?i)^TABLE\\s+`?([A-Za-z_][A-Za-z0-9_]*)`?\\s*\\((.*)\\)").matcher(line);
            if (!table.find()) continue;
            String tableName = table.group(1);
            for (String rawColumn : table.group(2).split(",")) {
                java.util.regex.Matcher column = Pattern.compile("\\s*`?([A-Za-z_][A-Za-z0-9_]*)`?\\s+").matcher(rawColumn);
                if (!column.find()) continue;
                String columnName = column.group(1);
                if (isScopedEntityField(tableName + "." + columnName)) result.add(new ScopedField(tableName, columnName));
            }
        }
        return result;
    }

    private static String quoteIdentifier(String identifier) {
        return "`" + identifier.replace("`", "") + "`";
    }

    private static String sqlLiteral(String value) {
        return "'" + String.valueOf(value).replace("'", "''") + "'";
    }

    private static String valuesForTable(String table, StructuredQueryContext context) {
        StringBuilder text = new StringBuilder();
        context.sampledValues().forEach((field, values) -> {
            if (field.startsWith(table + ".")) values.forEach(value -> text.append(' ').append(value));
        });
        context.referenceRows().getOrDefault(table, List.of()).forEach(row ->
                row.values().forEach(value -> text.append(' ').append(value)));
        return text.toString();
    }

    private static boolean containsAny(String text, String... values) {
        for (String value : values) {
            if (text.contains(value)) return true;
        }
        return false;
    }

    private static String summary(SchemaExplorerService.QueryResult result) {
        return "查询返回 " + result.rows().size() + " 行，列：" + String.join(", ", result.columns());
    }

    private record MetricDiagnostic(String code, String message, Map<String, Object> metadata) {}

    private record PreflightResult(Optional<String> sql, int repairCount,
                                   List<String> usedRelationships,
                                   List<String> relationshipWarnings) {}

    private record ScopedField(String table, String column) {}
}
