package com.agent.management.rag.graph;

import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;

@Component
public class GraphEntityNeighborhoodQueryBuilder {
    private static final int MAX_RELATIONSHIPS = 4;

    public Optional<String> build(GraphSchemaSnapshot schema,
                                  GraphEntityGroundingService.GroundingResult grounding,
                                  int maxRows) {
        if (grounding.matches().isEmpty()) return Optional.empty();
        GraphEntityGroundingService.EntityMatch entity = grounding.matches().get(0);
        String entityLabel = entity.labels().stream()
                .filter(label -> schema.node(label)
                        .map(node -> node.properties().containsKey(entity.property())).orElse(false))
                .findFirst().orElse(null);
        if (entityLabel == null) return Optional.empty();

        List<Neighbor> neighbors = neighbors(schema, new LinkedHashSet<>(entity.labels()));
        StringBuilder cypher = new StringBuilder("MATCH (n:")
                .append(identifier(entityLabel)).append(" {")
                .append(identifier(entity.property())).append(": ")
                .append(literal(entity.mention())).append("})\n");
        List<String> returns = new ArrayList<>(List.of("n"));
        for (int index = 0; index < neighbors.size(); index++) {
            Neighbor neighbor = neighbors.get(index);
            String pathVariable = "p" + (index + 1);
            String nodeVariable = "m" + (index + 1);
            cypher.append("OPTIONAL MATCH ").append(pathVariable).append("=");
            if (neighbor.outgoing()) {
                cypher.append("(n)-[:")
                        .append(identifier(neighbor.type())).append("]->(")
                        .append(nodeVariable).append(":").append(identifier(neighbor.otherLabel())).append(")\n");
            } else {
                cypher.append("(").append(nodeVariable).append(":")
                        .append(identifier(neighbor.otherLabel())).append(")-[:")
                        .append(identifier(neighbor.type())).append("]->(n)\n");
            }
            returns.add(pathVariable);
        }
        cypher.append("RETURN ").append(String.join(", ", returns))
                .append(" LIMIT ").append(Math.max(1, maxRows));
        return Optional.of(cypher.toString());
    }

    private static List<Neighbor> neighbors(GraphSchemaSnapshot schema, Set<String> entityLabels) {
        LinkedHashSet<Neighbor> result = new LinkedHashSet<>();
        for (GraphSchemaSnapshot.RelationshipSchema relationship : schema.relationships()) {
            for (GraphSchemaSnapshot.RelationshipEndpoint endpoint : relationship.endpoints()) {
                if (entityLabels.contains(endpoint.startLabel())) {
                    result.add(new Neighbor(relationship.type(), endpoint.endLabel(), true));
                } else if (entityLabels.contains(endpoint.endLabel())) {
                    result.add(new Neighbor(relationship.type(), endpoint.startLabel(), false));
                }
                if (result.size() >= MAX_RELATIONSHIPS) return new ArrayList<>(result);
            }
        }
        return new ArrayList<>(result);
    }

    private static String identifier(String value) {
        return "`" + value.replace("`", "``") + "`";
    }

    private static String literal(String value) {
        return "'" + value.replace("\\", "\\\\").replace("'", "\\'") + "'";
    }

    private record Neighbor(String type, String otherLabel, boolean outgoing) {}
}
