package com.agent.management.rag.graph;

import com.agent.management.rag.api.RagRetriever;
import com.agent.management.rag.memory.RagExampleMemoryService;
import com.agent.management.rag.model.*;
import com.agent.management.service.GraphSourceService;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.*;

@Service
@RequiredArgsConstructor
public class GraphRagRetriever implements RagRetriever {
    private final ExplicitCypherGenerationService explicit;
    private final Neo4jGraphRagCypherGenerationService generated;
    private final GraphSourceService graphs;
    private final RagExampleMemoryService exampleMemory;

    public RagSourceType sourceType() { return RagSourceType.GRAPH; }

    public RagRetrievalResult retrieve(RagQuery query) {
        RagRetrievalResult out = new RagRetrievalResult();
        out.setSourceType(sourceType());
        try {
            Long id = sourceId(query);
            Optional<String> cypher = explicit.generateCypher(query, id);
            boolean auto = false;
            int repairCount = 0;
            Map<String,Object> generationMetadata = Map.of();
            if (cypher.isEmpty()) {
                auto = true;
                try {
                    Optional<Neo4jGraphRagCypherGenerationService.GeneratedCypher> generation = generated.generate(query, id);
                    cypher = generation.map(Neo4jGraphRagCypherGenerationService.GeneratedCypher::cypher);
                    generationMetadata = generation.map(Neo4jGraphRagCypherGenerationService.GeneratedCypher::contextMetadata)
                            .orElse(Map.of());
                    repairCount = generation.filter(Neo4jGraphRagCypherGenerationService.GeneratedCypher::repaired)
                            .isPresent() ? 1 : 0;
                } catch (Exception error) {
                    out.getDiagnostics().put("error", "Cypher generation failed: " + error.getMessage());
                    return out;
                }
            }
            if (cypher.isEmpty()) {
                out.getDiagnostics().put("warning", "No Cypher was generated. The generic graph planner did not find a supported intent, and Text2Cypher returned empty content.");
                out.getDiagnostics().put("suggestions", List.of(
                        "Use an explicit entity id/name and ask for its related nodes, status, evidence, deployment, or capabilities.",
                        "Ask count/ranking questions only when the graph has related entities or score/rank/priority properties.",
                        "For capability matching, ask about capabilities that exist as Capability node property values.",
                        "Provide explicit Cypher in filters.cypher for unsupported graph question types."));
                out.getDiagnostics().put("supportedGenerationModes", List.of(
                        "ENTITY_NEIGHBORHOOD", "COUNT_RELATED", "PROPERTY_FILTERED_EVIDENCE",
                        "RELATION_RANKING", "CAPABILITY_MATCHING", "LLM_TEXT2CYPHER"));
                return out;
            }
            if (!auto) generated.validateExplicit(query, id, cypher.get());
            if (auto) {
                PreflightResult preflight = preflightAndRepair(query, id, cypher.get(), out, repairCount == 0);
                cypher = preflight.cypher();
                repairCount += preflight.repairCount();
                if (cypher.isEmpty()) return out;
            }
            Map<String, Object> result = graphs.executeQuery(id, cypher.get());
            if (auto && repairCount == 0 && !hasResults(result)) {
                try {
                    Optional<String> repaired = generated.repair(query, id, cypher.get(),
                            "query executed successfully but returned zero results; entity names may be abbreviated, use CONTAINS on an authorized name/title property");
                    if (repaired != null && repaired.isPresent() && !repaired.get().equals(cypher.get())) {
                        graphs.explainQuery(id, repaired.get());
                        Map<String,Object> retried = graphs.executeQuery(id, repaired.get());
                        if (hasResults(retried)) {
                            cypher = repaired; result = retried;
                            repairCount = 1;
                            out.getDiagnostics().put("repair", "zero-result Cypher was repaired once with fuzzy entity matching");
                        } else {
                            out.getDiagnostics().put("warning", "generated Cypher returned zero results; repair also returned zero results");
                        }
                    }
                } catch (Exception repairError) {
                    out.getDiagnostics().put("warning", "generated Cypher returned zero results; repair failed: "
                            + repairError.getMessage());
                }
            }
            List<?> nodes = list(result.get("nodes"));
            List<?> edges = list(result.get("edges"));
            List<?> records = list(result.get("records"));
            RagEvidence evidence = new RagEvidence();
            evidence.setId("graph-" + UUID.randomUUID());
            evidence.setSourceType(sourceType());
            evidence.setEvidenceType("GRAPH_RESULT");
            evidence.setTitle(graphs.get(id).getName());
            evidence.setSourceName(evidence.getTitle());
            evidence.setContent("图查询返回 " + nodes.size() + " 个节点、" + edges.size()
                    + " 条关系、" + records.size() + " 行记录");
            evidence.setScore(1.0);
            RagSourceRef ref = new RagSourceRef();
            ref.setSourceId(String.valueOf(id));
            ref.setSourceName(evidence.getTitle());
            ref.setLocator(Map.of("graphSourceId", id));
            evidence.setSourceRef(ref);
            evidence.setPayload(Map.of("nodes", nodes, "edges", edges, "records", records, "paths", List.of()));
            Map<String, Object> metadata = new LinkedHashMap<>();
            metadata.put("cypher", cypher.get());
            metadata.put(auto ? "generatedCypher" : "explicitCypher", cypher.get());
            metadata.put("nodeCount", nodes.size());
            metadata.put("edgeCount", edges.size());
            metadata.put("recordCount", records.size());
            metadata.put("durationMs", result.getOrDefault("durationMs", 0));
            metadata.put("warnings", Boolean.TRUE.equals(result.get("truncated")) ? List.of("result truncated") : List.of());
            metadata.put("repaired", repairCount > 0);
            metadata.put("repairCount", repairCount);
            metadata.putAll(generationMetadata);
            evidence.setMetadata(metadata);
            out.getEvidences().add(evidence);
            if (auto && autoLearn(query) && hasResults(result))
                exampleMemory.recordSuccess(sourceType(), String.valueOf(id), query.getQuery(), cypher.get());
        } catch (Exception error) {
            out.getDiagnostics().put("error", error.getMessage());
        }
        return out;
    }

    private PreflightResult preflightAndRepair(RagQuery query, Long sourceId, String cypher,
                                               RagRetrievalResult out, boolean allowRepair) {
        try {
            graphs.explainQuery(sourceId, cypher);
            return new PreflightResult(Optional.of(cypher), 0);
        } catch (Exception explainError) {
            if (!allowRepair) {
                out.getDiagnostics().put("error", "Cypher EXPLAIN failed after repair budget was exhausted: "
                        + explainError.getMessage());
                return new PreflightResult(Optional.empty(), 0);
            }
            try {
                Optional<String> repaired = generated.repair(query, sourceId, cypher, explainError.getMessage());
                if (repaired.isEmpty()) throw new IllegalArgumentException("repair returned empty Cypher");
                graphs.explainQuery(sourceId, repaired.get());
                out.getDiagnostics().put("repair", "generated Cypher failed EXPLAIN and was repaired once");
                return new PreflightResult(repaired, 1);
            } catch (Exception repairError) {
                out.getDiagnostics().put("error", "Cypher preflight/repair failed: " + repairError.getMessage());
                return new PreflightResult(Optional.empty(), 1);
            }
        }
    }

    private static boolean autoLearn(RagQuery query) {
        return query.getFilters() == null || !Boolean.FALSE.equals(query.getFilters().get("autoLearnExamples"));
    }

    private static Long sourceId(RagQuery query) {
        if (query.getSourceIds() == null || query.getSourceIds().isEmpty()) throw new IllegalArgumentException("sourceIds must contain graphSourceId");
        if (query.getSourceIds().size() != 1) throw new IllegalArgumentException("Text2Cypher supports exactly one graphSourceId per request");
        return Long.valueOf(query.getSourceIds().get(0));
    }

    private static List<?> list(Object value) { return value instanceof List<?> list ? list : List.of(); }
    private static boolean hasResults(Map<String, Object> result) {
        return !list(result.get("nodes")).isEmpty() || !list(result.get("edges")).isEmpty()
                || !list(result.get("records")).isEmpty();
    }

    private record PreflightResult(Optional<String> cypher, int repairCount) {}
}
