package com.agent.management.rag.memory;

import com.agent.management.model.entity.RagQueryExample;
import com.agent.management.rag.capability.RagTextSimilarity;
import com.agent.management.rag.capability.RagSemanticCatalogService;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.repository.RagQueryExampleRepository;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.*;

@Service
@RequiredArgsConstructor
public class RagExampleMemoryService {
    private final RagQueryExampleRepository repository;
    private final RagSemanticCatalogService semanticCatalog;
    private final com.agent.management.rag.capability.RagCapabilityProfileService profiles;

    public List<Map<String, Object>> findSimilar(RagSourceType type, String sourceId, String question, int limit) {
        String version = profiles.get(type, sourceId).version();
        List<RagQueryExample> candidates = repository.findBySourceTypeAndSourceIdAndVerifiedTrue(type, sourceId).stream()
                .filter(item -> {
                    if (item.getSchemaVersion() == null) { item.setSchemaVersion(version); repository.save(item); return true; }
                    if (version.equals(item.getSchemaVersion())) return true;
                    item.setVerified(false); repository.save(item); return false;
                }).toList();
        Map<String,String> texts = new LinkedHashMap<>();
        for (int i = 0; i < candidates.size(); i++) texts.put(String.valueOf(i), candidates.get(i).getQuestion());
        List<String> semanticOrder = semanticCatalog.rankTexts(question, texts, limit);
        java.util.stream.Stream<RagQueryExample> ordered = semanticOrder.isEmpty() ? candidates.stream()
                .sorted(Comparator.comparingDouble((RagQueryExample item) ->
                        RagTextSimilarity.score(question, item.getQuestion())).reversed())
                : semanticOrder.stream().map(index -> candidates.get(Integer.parseInt(index)));
        return ordered.limit(Math.max(0, limit))
                .map(item -> Map.<String, Object>of("question", item.getQuestion(), "query", item.getQueryText()))
                .toList();
    }

    public void recordSuccess(RagSourceType type, String sourceId, String question, String queryText) {
        if (corrupted(question) || corrupted(queryText)) return;
        repository.findFirstBySourceTypeAndSourceIdAndQuestionAndQueryText(type, sourceId, question, queryText)
                .ifPresentOrElse(existing -> {
                    existing.setSuccessCount(existing.getSuccessCount() + 1);
                    repository.save(existing);
                }, () -> {
                    RagQueryExample example = new RagQueryExample();
                    example.setSourceType(type); example.setSourceId(sourceId); example.setQuestion(question);
                    example.setQueryText(queryText); example.setVerified(false);
                    example.setSchemaVersion(profiles.get(type, sourceId).version()); repository.save(example);
                });
    }

    public List<RagQueryExample> list(RagSourceType type, String sourceId) {
        return repository.findBySourceTypeAndSourceId(type, sourceId);
    }

    public RagQueryExample setVerified(Long id, boolean verified) {
        RagQueryExample example = repository.findById(id)
                .orElseThrow(() -> new IllegalArgumentException("query example does not exist: " + id));
        example.setVerified(verified);
        return repository.save(example);
    }

    public void delete(Long id) {
        if (!repository.existsById(id)) throw new IllegalArgumentException("query example does not exist: " + id);
        repository.deleteById(id);
    }

    public List<RagQueryExample> addCandidates(RagSourceType type, String sourceId, List<Map<String,Object>> candidates) {
        for (Map<String,Object> candidate : candidates == null ? List.<Map<String,Object>>of() : candidates) {
            String question = String.valueOf(candidate.getOrDefault("question", "")).trim();
            Object rawQuery = candidate.containsKey("sql") ? candidate.get("sql")
                    : candidate.containsKey("cypher") ? candidate.get("cypher") : candidate.get("query");
            String query = rawQuery == null ? "" : String.valueOf(rawQuery).trim();
            if (!question.isEmpty() && !query.isEmpty() && readOnlyCandidate(query)) recordSuccess(type, sourceId, question, query);
        }
        return list(type, sourceId);
    }

    private static boolean readOnlyCandidate(String query) {
        String normalized = query.toUpperCase(Locale.ROOT);
        return java.util.regex.Pattern.compile("\\b(INSERT|UPDATE|DELETE|DROP|ALTER|TRUNCATE|CREATE|MERGE|SET|REMOVE|LOAD\\s+CSV)\\b")
                .matcher(normalized).find() == false;
    }

    private static boolean corrupted(String value) {
        if (value == null || value.isBlank() || value.indexOf('\uFFFD') >= 0) return true;
        long questionMarks = value.chars().filter(ch -> ch == '?').count();
        return value.contains("????") || questionMarks >= Math.max(4, value.length() / 3);
    }
}
