package com.agent.management.rag.graph;

import java.util.*;
import java.util.stream.Collectors;

public record GraphSchemaSnapshot(List<NodeSchema> nodes, List<RelationshipSchema> relationships) {
    public GraphSchemaSnapshot {
        nodes = List.copyOf(nodes == null ? List.of() : nodes);
        relationships = List.copyOf(relationships == null ? List.of() : relationships);
    }

    public record NodeSchema(String label, Map<String, List<String>> properties) {
        public NodeSchema {
            properties = immutableProperties(properties);
        }
    }

    public record RelationshipEndpoint(String startLabel, String endLabel) {}

    public record RelationshipSchema(String type, List<String> startLabels, List<String> endLabels,
                                     Map<String, List<String>> properties,
                                     List<RelationshipEndpoint> endpoints) {
        public RelationshipSchema(String type, List<String> startLabels, List<String> endLabels,
                                  Map<String, List<String>> properties) {
            this(type, startLabels, endLabels, properties, cartesianEndpoints(startLabels, endLabels));
        }

        public RelationshipSchema {
            startLabels = List.copyOf(startLabels == null ? List.of() : startLabels);
            endLabels = List.copyOf(endLabels == null ? List.of() : endLabels);
            properties = immutableProperties(properties);
            endpoints = List.copyOf(endpoints == null ? List.of() : endpoints);
        }
    }

    public Optional<NodeSchema> node(String label) {
        return nodes.stream().filter(node -> node.label().equals(label)).findFirst();
    }

    public List<String> labelNames() {
        return nodes.stream().map(NodeSchema::label).toList();
    }

    public List<String> relationshipTypeNames() {
        return relationships.stream().map(RelationshipSchema::type).distinct().toList();
    }

    public GraphSchemaSnapshot filter(List<String> allowedLabels, List<String> allowedRelationships,
                                      Map<String, List<String>> allowedProperties) {
        Set<String> labelFilter = normalizedSet(allowedLabels);
        Set<String> relationshipFilter = normalizedSet(allowedRelationships);
        List<NodeSchema> filteredNodes = nodes.stream()
                .filter(node -> labelFilter.isEmpty() || labelFilter.contains(node.label()))
                .map(node -> new NodeSchema(node.label(), filterProperties(node.label(), node.properties(), allowedProperties)))
                .toList();
        Set<String> retainedLabels = filteredNodes.stream().map(NodeSchema::label).collect(Collectors.toSet());
        List<RelationshipSchema> filteredRelationships = relationships.stream()
                .filter(rel -> relationshipFilter.isEmpty() || relationshipFilter.contains(rel.type()))
                .map(rel -> Map.entry(rel, rel.endpoints().stream()
                        .filter(endpoint -> retainedLabels.contains(endpoint.startLabel())
                                && retainedLabels.contains(endpoint.endLabel())).toList()))
                .filter(entry -> !entry.getValue().isEmpty())
                .map(entry -> new RelationshipSchema(entry.getKey().type(),
                        entry.getValue().stream().map(RelationshipEndpoint::startLabel).distinct().toList(),
                        entry.getValue().stream().map(RelationshipEndpoint::endLabel).distinct().toList(),
                        filterProperties(entry.getKey().type(), entry.getKey().properties(), allowedProperties),
                        entry.getValue()))
                .toList();
        return new GraphSchemaSnapshot(filteredNodes, filteredRelationships);
    }

    public String toPromptText() {
        StringBuilder text = new StringBuilder("Nodes:\n");
        for (NodeSchema node : nodes) {
            text.append("(:").append(node.label()).append(" {")
                    .append(formatProperties(node.properties())).append("})\n");
        }
        text.append("Relationships:\n");
        for (RelationshipSchema rel : relationships) {
            for (RelationshipEndpoint endpoint : rel.endpoints()) {
                text.append("(:").append(endpoint.startLabel()).append(")-[:").append(rel.type());
                if (!rel.properties().isEmpty()) text.append(" {").append(formatProperties(rel.properties())).append("}");
                text.append("]->(:").append(endpoint.endLabel()).append(")\n");
            }
        }
        return text.toString();
    }

    private static Map<String, List<String>> filterProperties(String owner, Map<String, List<String>> actual,
                                                               Map<String, List<String>> allowed) {
        if (allowed == null || allowed.isEmpty()) return actual;
        List<String> configured = allowed.get(owner);
        if (configured == null) configured = allowed.get("*");
        if (configured == null) return Map.of();
        Set<String> names = new HashSet<>(configured);
        return actual.entrySet().stream().filter(entry -> names.contains(entry.getKey()))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, (a, b) -> a, LinkedHashMap::new));
    }

    private static Set<String> normalizedSet(List<String> values) {
        return values == null ? Set.of() : values.stream().filter(Objects::nonNull)
                .collect(Collectors.toCollection(LinkedHashSet::new));
    }

    private static Map<String, List<String>> immutableProperties(Map<String, List<String>> source) {
        if (source == null || source.isEmpty()) return Map.of();
        Map<String, List<String>> result = new LinkedHashMap<>();
        source.forEach((key, value) -> result.put(key, List.copyOf(value == null ? List.of() : value)));
        return Collections.unmodifiableMap(result);
    }

    private static List<RelationshipEndpoint> cartesianEndpoints(List<String> starts, List<String> ends) {
        List<RelationshipEndpoint> result = new ArrayList<>();
        for (String start : starts == null ? List.<String>of() : starts) {
            for (String end : ends == null ? List.<String>of() : ends) {
                result.add(new RelationshipEndpoint(start, end));
            }
        }
        return result;
    }

    private static String formatProperties(Map<String, List<String>> properties) {
        return properties.entrySet().stream()
                .map(entry -> entry.getKey() + ": " + String.join("|", entry.getValue()))
                .collect(Collectors.joining(", "));
    }
}
