package com.agent.management.rag.graph;

import com.agent.management.rag.bridge.RagAiBridgeClient;
import com.agent.management.rag.capability.RagCapabilityProfileService;
import com.agent.management.rag.capability.RagSchemaLinker;
import com.agent.management.rag.capability.RagCapabilityProfile;
import com.agent.management.rag.capability.GraphBusinessSubgraphSelector;
import com.agent.management.rag.capability.RagEntityMentionExtractor;
import com.agent.management.rag.memory.RagExampleMemoryService;
import com.agent.management.rag.model.RagQuery;
import com.agent.management.rag.model.RagSourceType;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.util.*;

@Service
@RequiredArgsConstructor
public class Neo4jGraphRagCypherGenerationService implements CypherGenerationService {
    private final RagAiBridgeClient bridge;
    private final RagCapabilityProfileService profiles;
    private final RagSchemaLinker linker;
    private final GraphBusinessSubgraphSelector subgraphs;
    private final GraphSchemaValidator validator;
    private final RagExampleMemoryService examples;
    private final RagEntityMentionExtractor entities;
    private final GraphEntityGroundingService groundingService;
    private final GraphEntityNeighborhoodQueryBuilder neighborhoodQueries;
    private final GraphQueryIntentPlanner intentPlanner;
    private final com.agent.management.rag.capability.GraphGovernanceConfigService governance;

    @Override
    public Optional<String> generateCypher(RagQuery query, Long graphSourceId) {
        return generate(query, graphSourceId).map(GeneratedCypher::cypher);
    }

    public Optional<GeneratedCypher> generate(RagQuery query, Long graphSourceId) {
        query = governance.apply(query, graphSourceId);
        if (!allowed(query)) return Optional.empty();
        List<Map<String,Object>> promptExamples = new ArrayList<>(maps(query, "examples"));
        promptExamples.addAll(examples.findSimilar(RagSourceType.GRAPH, String.valueOf(graphSourceId), query.getQuery(), 3));
        SelectionContext selection = selection(query, graphSourceId, promptExamples);
        GraphSchemaSnapshot selected = selection.schema();
        if (selected.nodes().isEmpty()) throw new IllegalArgumentException("no relevant authorized graph schema is available for Text-to-Cypher");
        int maxDepth = integer(query, "maxDepth", 3);
        Optional<GeneratedCypher> plannedQuery = plannedQuery(query, graphSourceId, selection, maxDepth);
        if (plannedQuery.isPresent()) return plannedQuery;
        Optional<GeneratedCypher> groundedQuery = neighborhoodQuery(selection, maxDepth);
        if (groundedQuery.isPresent()) return groundedQuery;

        Map<String,Object> request = new LinkedHashMap<>();
        request.put("query",query.getQuery()); request.put("graphSourceId",graphSourceId);
        request.put("schema",selected.toPromptText()); request.put("examples",normalizeExamples(promptExamples));
        request.put("businessRules",generationContext(query, selection)); request.put("allowedLabels",selected.labelNames());
        request.put("allowedRelationships",selected.relationshipTypeNames()); request.put("allowedProperties",allowedProperties(selected));
        request.put("maxDepth",maxDepth);
        request.put("entityMentions", entityMentions(query, selection));
        Object cypher = bridge.textToCypher(request).get("cypher");
        if(cypher==null||String.valueOf(cypher).isBlank()) return Optional.empty();
        String generated=String.valueOf(cypher);
        try {
            validator.validate(generated,selected,maxDepth);
            return Optional.of(new GeneratedCypher(generated, false, selection.metadata()));
        } catch (IllegalArgumentException validationError) {
            Object fixed=bridge.repair(Map.of("language","CYPHER","question",query.getQuery(),"query",generated,
                    "error",validationError.getMessage(),"schemaText",selected.toPromptText(),"maxRows",50,"maxDepth",maxDepth),
                    Duration.ofSeconds(30)).get("query");
            if(fixed==null||String.valueOf(fixed).isBlank()) throw validationError;
            String repaired=String.valueOf(fixed);validator.validate(repaired,selected,maxDepth);
            return Optional.of(new GeneratedCypher(repaired, true, selection.metadata()));
        }
    }

    private Optional<GeneratedCypher> neighborhoodQuery(SelectionContext selection, int maxDepth) {
        return neighborhoodQueries.build(selection.schema(), selection.grounding(), 50).map(cypher -> {
            validator.validate(cypher, selection.schema(), maxDepth);
            Map<String,Object> metadata = new LinkedHashMap<>(selection.metadata());
            metadata.put("generationMode", "VERIFIED_ENTITY_NEIGHBORHOOD");
            return new GeneratedCypher(cypher, false, metadata);
        });
    }

    private Optional<GeneratedCypher> plannedQuery(RagQuery query, Long graphSourceId,
                                                   SelectionContext selection, int maxDepth) {
        return intentPlanner.plan(query, graphSourceId, selection.authorizedSchema(), selection.grounding(), maxDepth)
                .map(planned -> {
                    Map<String,Object> metadata = new LinkedHashMap<>(selection.metadata(selection.authorizedSchema()));
                    metadata.put("generationMode", planned.generationMode());
                    metadata.put("intentPlan", planned.intentPlan());
                    metadata.put("groundingConfidence", planned.groundingConfidence());
                    metadata.put("fallbackReason", "generic graph intent planner");
                    return new GeneratedCypher(planned.cypher(), false, metadata);
                });
    }

    public Optional<String> repair(RagQuery query,Long graphSourceId,String failedCypher,String error){
        query = governance.apply(query, graphSourceId);
        List<Map<String,Object>> promptExamples = new ArrayList<>(maps(query, "examples"));
        promptExamples.addAll(examples.findSimilar(RagSourceType.GRAPH, String.valueOf(graphSourceId), query.getQuery(), 3));
        GraphSchemaSnapshot schema = selection(query, graphSourceId, promptExamples).schema();
        int maxDepth=integer(query,"maxDepth",3);
        Object fixed=bridge.repair(Map.of("language","CYPHER","question",query.getQuery(),"query",failedCypher,
                "error",error,"schemaText",schema.toPromptText(),"maxRows",50,"maxDepth",maxDepth),
                Duration.ofSeconds(30)).get("query");
        if(fixed==null||String.valueOf(fixed).isBlank())return Optional.empty();
        validator.validate(String.valueOf(fixed),schema,maxDepth);return Optional.of(String.valueOf(fixed));
    }

    public void validateExplicit(RagQuery query, Long graphSourceId, String cypher) {
        query = governance.apply(query, graphSourceId);
        GraphSchemaSnapshot authorized = profiles.get(RagSourceType.GRAPH, String.valueOf(graphSourceId)).graphSchema()
                .filter(stringList(query, "allowedLabels"), stringList(query, "allowedRelationships"), properties(query));
        validator.validate(cypher, authorized, integer(query, "maxDepth", 3));
    }

    static boolean allowed(RagQuery q){if(q.getFilters()==null)return false;return Boolean.TRUE.equals(q.getFilters().get("allowTextToCypher"))||"AUTO_GENERATE".equals(String.valueOf(q.getFilters().get("retrievalMode")));}
    private static int integer(RagQuery q,String key,int d){Object v=q.getFilters().get(key);return v instanceof Number n?n.intValue():d;}
    private static String string(RagQuery q,String key){Object value=q.getFilters().get(key);return value==null?"":String.valueOf(value);}
    private static String generationContext(RagQuery query, SelectionContext selection){String existing=string(query,"businessRules");List<String> rules=stringList(query,"generationRules");String context=rules.isEmpty()?existing:existing+(existing.isBlank()?"":"\n")+"Generation rules:\n- "+String.join("\n- ",rules);if(!selection.grounding().matches().isEmpty())context+=(context.isBlank()?"":"\n")+"Verified graph entities:\n- "+String.join("\n- ",selection.grounding().promptMentions());return context;}
    private static List<String> stringList(RagQuery q,String key){Object value=q.getFilters().get(key);return value instanceof List<?> list?list.stream().map(String::valueOf).toList():List.of();}
    private static List<Map<String,Object>> maps(RagQuery q,String key){Object value=q.getFilters().get(key);if(!(value instanceof List<?> list))return List.of();return list.stream().filter(Map.class::isInstance).map(item->{Map<?,?> raw=(Map<?,?>)item;Map<String,Object> map=new LinkedHashMap<>();raw.forEach((k,v)->map.put(String.valueOf(k),v));return map;}).toList();}
    private static List<Map<String,Object>> normalizeExamples(List<Map<String,Object>> examples){return examples.stream().map(item->item.containsKey("cypher")?item:Map.<String,Object>of("question",item.get("question"),"cypher",item.get("query"))).toList();}
    private static Map<String,List<String>> properties(RagQuery q){Object value=q.getFilters().get("allowedProperties");if(!(value instanceof Map<?,?> map))return Map.of();Map<String,List<String>> result=new LinkedHashMap<>();map.forEach((key,raw)->{if(raw instanceof List<?> list)result.put(String.valueOf(key),list.stream().map(String::valueOf).toList());});return result;}
    private static Map<String,List<String>> allowedProperties(GraphSchemaSnapshot schema){Map<String,List<String>> result=new LinkedHashMap<>();schema.nodes().forEach(node->result.put(node.label(),new ArrayList<>(node.properties().keySet())));schema.relationships().forEach(rel->result.put(rel.type(),new ArrayList<>(rel.properties().keySet())));return result;}

    private SelectionContext selection(RagQuery query, Long graphSourceId,
                                       List<Map<String,Object>> promptExamples) {
        RagCapabilityProfile profile = profiles.get(RagSourceType.GRAPH, String.valueOf(graphSourceId));
        GraphSchemaSnapshot authorized = profile.graphSchema().filter(stringList(query,"allowedLabels"),
                stringList(query,"allowedRelationships"), properties(query));
        GraphEntityGroundingService.GroundingResult grounding =
                groundingService.ground(graphSourceId, query.getQuery(), authorized);
        GraphSchemaSnapshot selected;
        if (!grounding.labels().isEmpty()) {
            selected = subgraphs.select(query.getQuery(), profile, authorized, grounding.labels(), 8, 10);
        } else {
            GraphSchemaSnapshot semanticSchema = authorized.nodes().size() <= 12 ? authorized
                    : subgraphs.select(query.getQuery(), profile, authorized, 8, 10);
            selected = linker.selectGraphSchema(query.getQuery(), semanticSchema, promptExamples);
        }
        if (selected.nodes().size() > 10 || selected.relationships().size() > 12) {
            selected = subgraphs.select(query.getQuery(), profile, selected, grounding.labels(), 8, 10);
        }
        if (selected.nodes().isEmpty()) {
            throw new IllegalArgumentException("no relevant authorized graph schema is available for Text-to-Cypher");
        }
        return new SelectionContext(selected, authorized, grounding);
    }

    private List<String> entityMentions(RagQuery query, SelectionContext selection) {
        LinkedHashSet<String> result = new LinkedHashSet<>(entities.extract(query.getQuery()));
        result.addAll(selection.grounding().promptMentions());
        return new ArrayList<>(result);
    }

    private record SelectionContext(GraphSchemaSnapshot schema,
                                    GraphSchemaSnapshot authorizedSchema,
                                    GraphEntityGroundingService.GroundingResult grounding) {
        Map<String,Object> metadata() {
            return metadata(schema);
        }
        Map<String,Object> metadata(GraphSchemaSnapshot metadataSchema) {
            Map<String,Object> result = new LinkedHashMap<>();
            result.put("selectedLabels", metadataSchema.labelNames());
            result.put("selectedRelationships", metadataSchema.relationshipTypeNames());
            result.put("entityMatches", grounding.metadata());
            result.put("schemaSelectionWarnings", grounding.warnings());
            return result;
        }
    }

    public record GeneratedCypher(String cypher, boolean repaired, Map<String,Object> contextMetadata) {
        public GeneratedCypher(String cypher, boolean repaired) { this(cypher, repaired, Map.of()); }
        public GeneratedCypher { contextMetadata = Map.copyOf(contextMetadata == null ? Map.of() : contextMetadata); }
    }
}
