package com.agent.management.rag.capability;

import com.agent.management.rag.graph.GraphSchemaSnapshot;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@Component
public class RagSchemaLinker {
    private static final Pattern LABEL = Pattern.compile(":`?([A-Za-z_]\\w*)`?");

    public List<String> selectSqlTables(String question, SqlSchemaSnapshot schema,
                                        List<String> policyWhitelist, Map<String, String> tableHints) {
        Set<String> policy = new HashSet<>(policyWhitelist == null ? List.of() : policyWhitelist);
        List<ScoredName> scores = new ArrayList<>();
        for (SqlSchemaSnapshot.TableSchema table : schema.tables()) {
            if (!policy.isEmpty() && !policy.contains(table.name())) continue;
            StringBuilder semantics = new StringBuilder(table.name()).append(' ').append(nullToEmpty(table.description()))
                    .append(' ').append(nullToEmpty(tableHints == null ? null : tableHints.get(table.name())));
            table.columns().forEach(column -> semantics.append(' ').append(column.name()).append(' ')
                    .append(nullToEmpty(column.description())));
            scores.add(new ScoredName(table.name(), RagTextSimilarity.score(question, semantics.toString())));
        }
        return select(scores);
    }

    public GraphSchemaSnapshot selectGraphSchema(String question, GraphSchemaSnapshot schema,
                                                  List<Map<String, Object>> examples) {
        Set<String> selectedNames = new LinkedHashSet<>();
        double bestExampleScore = 0;
        String bestQuery = null;
        for (Map<String, Object> example : examples == null ? List.<Map<String, Object>>of() : examples) {
            double score = RagTextSimilarity.score(question, String.valueOf(example.getOrDefault("question", "")));
            if (score > bestExampleScore) {
                bestExampleScore = score;
                Object query = example.containsKey("cypher") ? example.get("cypher") : example.get("query");
                bestQuery = query == null ? null : String.valueOf(query);
            }
        }
        if (bestExampleScore > 0.15 && bestQuery != null) {
            Matcher matcher = LABEL.matcher(bestQuery);
            while (matcher.find()) selectedNames.add(matcher.group(1));
        }
        if (selectedNames.isEmpty()) {
            for (String label : schema.labelNames()) {
                if (RagTextSimilarity.score(question, label) > 0) selectedNames.add(label);
            }
            for (String relationship : schema.relationshipTypeNames()) {
                if (RagTextSimilarity.score(question, relationship) > 0) selectedNames.add(relationship);
            }
        }
        if (selectedNames.isEmpty()) return schema;
        List<String> labels = schema.labelNames().stream().filter(selectedNames::contains).toList();
        List<String> relationships = schema.relationshipTypeNames().stream().filter(selectedNames::contains).toList();
        if (relationships.isEmpty()) {
            relationships = schema.relationships().stream()
                    .filter(rel -> rel.startLabels().stream().anyMatch(labels::contains)
                            && rel.endLabels().stream().anyMatch(labels::contains))
                    .map(GraphSchemaSnapshot.RelationshipSchema::type).distinct().toList();
        }
        Set<String> connected = new LinkedHashSet<>(labels);
        for (GraphSchemaSnapshot.RelationshipSchema rel : schema.relationships()) {
            if (relationships.contains(rel.type())) { connected.addAll(rel.startLabels()); connected.addAll(rel.endLabels()); }
        }
        return schema.filter(new ArrayList<>(connected), relationships, Map.of());
    }

    private static List<String> select(List<ScoredName> scores) {
        if (scores.isEmpty()) return List.of();
        double max = scores.stream().mapToDouble(ScoredName::score).max().orElse(0);
        if (max <= 0) return scores.stream().map(ScoredName::name).toList();
        double threshold = Math.max(0.08, max * 0.55);
        return scores.stream().filter(item -> item.score() >= threshold).map(ScoredName::name).toList();
    }

    private static String nullToEmpty(String value) { return value == null ? "" : value; }
    private record ScoredName(String name, double score) {}
}
