package com.agent.management.rag.capability;

import com.agent.management.rag.graph.GraphSchemaSnapshot;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.*;

@Service
@RequiredArgsConstructor
public class GraphBusinessSubgraphSelector {
    private final RagSemanticCatalogService semanticCatalog;

    public GraphSchemaSnapshot select(String intent, RagCapabilityProfile profile, GraphSchemaSnapshot allowed,
                                      int maxLabels, int maxRelationships) {
        return select(intent, profile, allowed, List.of(), maxLabels, maxRelationships);
    }

    public GraphSchemaSnapshot select(String intent, RagCapabilityProfile profile, GraphSchemaSnapshot allowed,
                                      List<String> preferredLabels, int maxLabels, int maxRelationships) {
        Set<String> allowedLabels = new LinkedHashSet<>(allowed.labelNames());
        Set<String> allowedRelationships = new LinkedHashSet<>(allowed.relationshipTypeNames());
        List<String> rankedLabels = semanticCatalog.rank(intent, profile, "label:", Math.max(4, maxLabels)).stream()
                .filter(allowedLabels::contains).toList();
        List<String> rankedRelationships = semanticCatalog.rank(intent, profile, "relationship:", Math.max(4, maxRelationships)).stream()
                .filter(allowedRelationships::contains).toList();

        LinkedHashSet<String> seeds = new LinkedHashSet<>();
        if (preferredLabels != null) preferredLabels.stream().filter(allowedLabels::contains).forEach(seeds::add);
        if (seeds.isEmpty()) rankedLabels.stream().filter(label -> seeds.size() < maxLabels).forEach(seeds::add);
        if (seeds.isEmpty()) {
            allowed.nodes().stream()
                    .sorted(Comparator.comparingDouble((GraphSchemaSnapshot.NodeSchema node) ->
                            RagTextSimilarity.score(intent, node.label() + " " + String.join(" ", node.properties().keySet()))).reversed())
                    .limit(Math.max(1, maxLabels / 2)).map(GraphSchemaSnapshot.NodeSchema::label).forEach(seeds::add);
        }

        Set<String> rankedRelSet = new HashSet<>(rankedRelationships);
        List<RelationshipScore> relationshipScores = allowed.relationships().stream().map(rel -> {
            boolean start = rel.startLabels().stream().anyMatch(seeds::contains);
            boolean end = rel.endLabels().stream().anyMatch(seeds::contains);
            int score = (rankedRelSet.contains(rel.type()) ? 4 : 0) + (start && end ? 3 : start || end ? 1 : 0);
            return new RelationshipScore(rel, score);
        }).filter(item -> item.score() > 0).sorted(Comparator.comparingInt(RelationshipScore::score).reversed())
                .limit(Math.max(1, maxRelationships)).toList();

        LinkedHashSet<String> selectedLabels = new LinkedHashSet<>(seeds.stream().limit(Math.max(1, maxLabels)).toList());
        LinkedHashSet<String> selectedRelationships = new LinkedHashSet<>();
        for (RelationshipScore item : relationshipScores) {
            if (selectedRelationships.size() >= maxRelationships) break;
            LinkedHashSet<String> required = new LinkedHashSet<>();
            item.relationship().endpoints().stream()
                    .filter(endpoint -> selectedLabels.contains(endpoint.startLabel())
                            || selectedLabels.contains(endpoint.endLabel()))
                    .forEach(endpoint -> {
                        required.add(endpoint.startLabel());
                        required.add(endpoint.endLabel());
                    });
            required.removeAll(selectedLabels);
            if (selectedLabels.size() + required.size() > maxLabels) continue;
            selectedRelationships.add(item.relationship().type());
            selectedLabels.addAll(required);
        }
        if (selectedRelationships.isEmpty() && !selectedLabels.isEmpty()) {
            allowed.relationships().stream()
                    .filter(rel -> rel.startLabels().stream().anyMatch(selectedLabels::contains)
                            && rel.endLabels().stream().anyMatch(selectedLabels::contains))
                    .limit(maxRelationships).map(GraphSchemaSnapshot.RelationshipSchema::type).forEach(selectedRelationships::add);
        }
        return allowed.filter(new ArrayList<>(selectedLabels), new ArrayList<>(selectedRelationships), Map.of());
    }

    private record RelationshipScore(GraphSchemaSnapshot.RelationshipSchema relationship, int score) {}
}
