package com.agent.management.rag.capability;

import com.agent.management.service.impl.EmbeddingBridgeClient;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

@Service
@RequiredArgsConstructor
public class RagSemanticCatalogService {
    private final EmbeddingBridgeClient embeddings;

    public List<String> rank(String question, RagCapabilityProfile profile, String prefix, int limit) {
        List<RagSemanticEntry> candidates = profile.semanticCatalog().stream()
                .filter(item -> item.key().startsWith(prefix) && !item.vector().isEmpty()).toList();
        if (candidates.isEmpty()) return List.of();
        try {
            List<List<Double>> vectors = embeddings.embed(List.of(question));
            if (vectors.isEmpty()) return List.of();
            List<Double> query = vectors.get(0);
            return candidates.stream().map(item -> new Score(item.key().substring(prefix.length()), cosine(query, item.vector())))
                    .filter(item -> item.score > 0.2).sorted(Comparator.comparingDouble(Score::score).reversed())
                    .limit(Math.max(1, limit)).map(Score::name).toList();
        } catch (Exception ignored) { return List.of(); }
    }

    public List<String> rankTexts(String question, Map<String,String> texts, int limit) {
        if (texts == null || texts.isEmpty()) return List.of();
        try {
            List<String> keys = List.copyOf(texts.keySet());
            List<String> inputs = new java.util.ArrayList<>();
            inputs.add(question); keys.forEach(key -> inputs.add(texts.get(key)));
            List<List<Double>> vectors = embeddings.embed(inputs);
            if (vectors.size() != inputs.size()) return List.of();
            List<Double> query = vectors.get(0);
            Map<String,Double> scores = new LinkedHashMap<>();
            for (int i = 0; i < keys.size(); i++) scores.put(keys.get(i), cosine(query, vectors.get(i + 1)));
            return scores.entrySet().stream().sorted(Map.Entry.<String,Double>comparingByValue().reversed())
                    .limit(Math.max(0, limit)).map(Map.Entry::getKey).toList();
        } catch (Exception ignored) { return List.of(); }
    }

    private static double cosine(List<Double> left, List<Double> right) {
        int size = Math.min(left.size(), right.size());
        double dot = 0, a = 0, b = 0;
        for (int i = 0; i < size; i++) { dot += left.get(i) * right.get(i); a += left.get(i) * left.get(i); b += right.get(i) * right.get(i); }
        return a == 0 || b == 0 ? 0 : dot / Math.sqrt(a * b);
    }
    private record Score(String name, double score) {}
}
