package com.agent.management.rag.structured;

import com.agent.management.rag.capability.RagEntityMentionExtractor;
import com.agent.management.rag.capability.RagTextSimilarity;
import com.agent.management.rag.capability.SqlSchemaSnapshot;
import com.agent.management.rag.model.RagQuery;
import com.agent.management.service.impl.EmbeddingBridgeClient;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;

@Service
@RequiredArgsConstructor
public class StructuredQueryContextService {
    private static final double LEXICAL_CONFIDENCE = 0.08;
    private static final double EMBEDDING_CONFIDENCE = 0.25;
    private static final int MAX_SEMANTIC_TABLES = 5;

    private final StructuredSchemaProfileService profiles;
    private final StructuredValueSampler sampler;
    private final RagEntityMentionExtractor entities;
    private final EmbeddingBridgeClient embeddings;

    public StructuredQueryContext build(RagQuery query, Long dataSourceId) {
        StructuredSchemaProfile profile = profiles.getProfile(dataSourceId);
        if (!"mysql".equalsIgnoreCase(profile.dialect())) {
            throw new IllegalArgumentException("automatic Text2SQL currently supports MySQL only");
        }
        if (!dataSourceId.equals(profile.dataSourceId())) {
            throw new IllegalStateException("structured schema profile belongs to another datasource");
        }
        SqlSchemaSnapshot schema = profile.schemaSnapshot();
        List<String> warnings = new ArrayList<>(profile.warnings());
        List<String> selectedTables = selectTables(query.getQuery(), schema, warnings);
        selectedTables = connectByRelationships(selectedTables, schema, profile.relationships());
        selectedTables = expandMetricContext(query.getQuery(), selectedTables, schema, profile.relationships());
        if (selectedTables.isEmpty()) {
            selectedTables = schema.tables().stream().map(SqlSchemaSnapshot.TableSchema::name).toList();
            warnings.add("table selection was empty; fell back to all visible schema tables");
        }
        String ddl = schema.toDdl(selectedTables);
        if (ddl.isBlank()) {
            selectedTables = schema.tables().stream().map(SqlSchemaSnapshot.TableSchema::name).toList();
            ddl = schema.toDdl(selectedTables);
            warnings.add("minimal DDL was empty; fell back to all visible schema tables");
        }

        List<String> selected = selectedTables;
        List<SqlSchemaSnapshot.ForeignKey> selectedForeignKeys = schema.foreignKeys().stream()
                .filter(key -> selected.contains(key.sourceTable()) && selected.contains(key.targetTable())).toList();
        List<StructuredRelationship> selectedRelationships = profile.relationships().stream()
                .filter(relation -> relation.confidence() >= 0.58)
                .filter(relation -> selected.contains(relation.sourceTable()) && selected.contains(relation.targetTable()))
                .toList();
        StructuredValueSampler.SampleResult samples = sampler.sample(
                dataSourceId, selectedTables, profile.columnSpecs());
        warnings.addAll(samples.warnings());
        return new StructuredQueryContext(
                dataSourceId,
                profile.dialect(),
                profile.schemaFingerprint(),
                profile.schemaFingerprint(),
                profile.profileVersion(),
                selectedTables,
                ddl,
                selectedForeignKeys,
                selectedRelationships,
                samples.sampledValues(),
                samples.referenceRows(),
                entities.extract(query.getQuery()),
                warnings);
    }

    List<String> selectTables(String question, SqlSchemaSnapshot schema, List<String> warnings) {
        List<TableScore> lexical = schema.tables().stream()
                .map(table -> new TableScore(table.name(), RagTextSimilarity.score(question, semantics(table))))
                .sorted(Comparator.comparingDouble(TableScore::score).reversed()).toList();
        double lexicalMax = lexical.stream().mapToDouble(TableScore::score).max().orElse(0);
        if (lexicalMax >= LEXICAL_CONFIDENCE) {
            double threshold = Math.max(LEXICAL_CONFIDENCE, lexicalMax * 0.55);
            return lexical.stream().filter(item -> item.score() >= threshold)
                    .limit(MAX_SEMANTIC_TABLES).map(TableScore::name).toList();
        }

        try {
            List<String> inputs = new ArrayList<>();
            inputs.add(question == null ? "" : question);
            schema.tables().stream().map(StructuredQueryContextService::semantics).forEach(inputs::add);
            List<List<Double>> vectors = embeddings.embed(inputs);
            if (vectors.size() == inputs.size()) {
                List<TableScore> semantic = new ArrayList<>();
                for (int i = 0; i < schema.tables().size(); i++) {
                    semantic.add(new TableScore(schema.tables().get(i).name(), cosine(vectors.get(0), vectors.get(i + 1))));
                }
                semantic.sort(Comparator.comparingDouble(TableScore::score).reversed());
                double max = semantic.isEmpty() ? 0 : semantic.get(0).score();
                if (max >= EMBEDDING_CONFIDENCE) {
                    double threshold = Math.max(EMBEDDING_CONFIDENCE, max - 0.08);
                    List<String> selected = semantic.stream().filter(item -> item.score() >= threshold)
                            .limit(MAX_SEMANTIC_TABLES).map(TableScore::name).toList();
                    if (!selected.isEmpty()) return selected;
                }
            }
        } catch (Exception error) {
            warnings.add("semantic table selection unavailable: " + shortMessage(error));
        }
        warnings.add("table selection confidence was low; fell back to all visible schema tables");
        return schema.tables().stream().map(SqlSchemaSnapshot.TableSchema::name).toList();
    }

    static List<String> connectByRelationships(List<String> selected, SqlSchemaSnapshot schema,
                                               List<StructuredRelationship> relationships) {
        List<StructuredRelationship> usable = relationships.stream()
                .filter(StructuredRelationship::usableForSelection).toList();
        LinkedHashSet<String> result = new LinkedHashSet<>(selected);
        for (int i = 0; i < selected.size(); i++) {
            for (int j = i + 1; j < selected.size(); j++) {
                result.addAll(shortestPath(selected.get(i), selected.get(j), usable));
            }
        }
        Set<String> schemaTables = schema.tables().stream().map(SqlSchemaSnapshot.TableSchema::name)
                .collect(java.util.stream.Collectors.toSet());
        return result.stream().filter(schemaTables::contains).toList();
    }

    static List<String> expandMetricContext(String question, List<String> selected, SqlSchemaSnapshot schema,
                                            List<StructuredRelationship> relationships) {
        if (!isMetricQuestion(question)) return selected;
        Set<String> schemaTables = schema.tables().stream().map(SqlSchemaSnapshot.TableSchema::name)
                .collect(java.util.stream.Collectors.toSet());
        Set<String> metricTables = schema.tables().stream()
                .filter(StructuredQueryContextService::hasMetricColumn)
                .map(SqlSchemaSnapshot.TableSchema::name)
                .collect(java.util.stream.Collectors.toCollection(LinkedHashSet::new));
        LinkedHashSet<String> result = new LinkedHashSet<>(selected);
        relationships.stream()
                .filter(StructuredRelationship::usableForSelection)
                .filter(relation -> metricTables.contains(relation.sourceTable())
                        || metricTables.contains(relation.targetTable()))
                .forEach(relation -> {
                    result.add(relation.sourceTable());
                    result.add(relation.targetTable());
                });
        return result.stream().filter(schemaTables::contains).toList();
    }

    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 boolean hasMetricColumn(SqlSchemaSnapshot.TableSchema table) {
        return table.columns().stream().anyMatch(column -> isMetricColumn(column.name()));
    }

    private static boolean isMetricColumn(String column) {
        String name = column == null ? "" : column.toLowerCase(Locale.ROOT);
        return name.equals("score") || name.endsWith("_score") || name.contains("rank")
                || name.contains("priority") || name.contains("rating")
                || name.equals("weight") || name.equals("level");
    }

    private static List<String> shortestPath(String start, String target,
                                             List<StructuredRelationship> relationships) {
        Map<String, Set<String>> graph = new HashMap<>();
        for (StructuredRelationship relation : relationships) {
            graph.computeIfAbsent(relation.sourceTable(), ignored -> new LinkedHashSet<>()).add(relation.targetTable());
            graph.computeIfAbsent(relation.targetTable(), ignored -> new LinkedHashSet<>()).add(relation.sourceTable());
        }
        ArrayDeque<List<String>> queue = new ArrayDeque<>();
        queue.add(List.of(start));
        Set<String> visited = new HashSet<>();
        while (!queue.isEmpty()) {
            List<String> path = queue.removeFirst();
            String current = path.get(path.size() - 1);
            if (!visited.add(current)) continue;
            if (current.equals(target)) return path;
            for (String next : graph.getOrDefault(current, Set.of())) {
                List<String> extended = new ArrayList<>(path); extended.add(next); queue.addLast(extended);
            }
        }
        return List.of();
    }

    private static String semantics(SqlSchemaSnapshot.TableSchema table) {
        StringBuilder text = new StringBuilder(table.name().replace('_', ' ')).append(' ')
                .append(table.description() == null ? "" : table.description());
        table.columns().forEach(column -> text.append(' ').append(column.name().replace('_', ' '))
                .append(' ').append(column.description() == null ? "" : column.description()));
        return text.toString();
    }

    private static double cosine(List<Double> left, List<Double> right) {
        if (left == null || right == null || left.size() != right.size() || left.isEmpty()) return 0;
        double dot = 0, leftNorm = 0, rightNorm = 0;
        for (int i = 0; i < left.size(); i++) {
            double a = left.get(i), b = right.get(i);
            dot += a * b; leftNorm += a * a; rightNorm += b * b;
        }
        return leftNorm == 0 || rightNorm == 0 ? 0 : dot / Math.sqrt(leftNorm * rightNorm);
    }

    private static String shortMessage(Exception error) {
        String message = error.getMessage();
        return message == null ? error.getClass().getSimpleName() : message.substring(0, Math.min(200, message.length()));
    }

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

    private record TableScore(String name, double score) {}
}
