package com.agent.management.rag.memory;

import com.agent.management.rag.capability.RagCapabilityProfileService;
import com.agent.management.rag.graph.GraphSchemaValidator;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.service.GraphSourceService;
import com.agent.management.service.SchemaExplorerService;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.*;

@Service
@RequiredArgsConstructor
public class RagExampleCandidateService {
    private final RagExampleMemoryService memory;
    private final RagCapabilityProfileService profiles;
    private final GraphSchemaValidator graphValidator;
    private final GraphSourceService graphSources;
    private final SchemaExplorerService sqlExplorer;

    public Map<String,Object> importValidated(RagSourceType type, String sourceId, List<Map<String,Object>> candidates) {
        int accepted = 0;
        List<Map<String,String>> rejected = new ArrayList<>();
        for (Map<String,Object> candidate : candidates == null ? List.<Map<String,Object>>of() : candidates) {
            String question = String.valueOf(candidate.getOrDefault("question", ""));
            try {
                String query = query(candidate);
                preflight(type, sourceId, query);
                memory.addCandidates(type, sourceId, List.of(candidate));
                accepted++;
            } catch (Exception error) {
                rejected.add(Map.of("question", question, "reason", shortMessage(error)));
            }
        }
        Map<String,Object> result = new LinkedHashMap<>();
        result.put("accepted", accepted); result.put("rejected", rejected);
        result.put("examples", memory.list(type, sourceId));
        return result;
    }

    private void preflight(RagSourceType type, String sourceId, String query) {
        if (query.isBlank()) throw new IllegalArgumentException("candidate query is empty");
        Long id = Long.valueOf(sourceId);
        if (type == RagSourceType.GRAPH) {
            var schema = profiles.get(type, sourceId).graphSchema();
            graphValidator.validate(query, schema, 5);
            graphSources.explainQuery(id, query);
        } else if (type == RagSourceType.STRUCTURED_DATA) {
            sqlExplorer.explainQuery(id, query);
        } else throw new IllegalArgumentException("Few-shot candidate is unsupported for " + type);
    }

    private static String query(Map<String,Object> candidate) {
        Object value = candidate.containsKey("sql") ? candidate.get("sql")
                : candidate.containsKey("cypher") ? candidate.get("cypher") : candidate.get("query");
        return value == null ? "" : String.valueOf(value).trim();
    }

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