package com.agent.management.rag.graph;

import com.agent.management.rag.capability.RagEntityMentionExtractor;
import com.agent.management.service.GraphSourceService;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.regex.Pattern;

@Service
@RequiredArgsConstructor
public class GraphEntityGroundingService {
    private static final Pattern IDENTIFIER = Pattern.compile("(?:[A-Z][A-Z0-9]*_)+(?:[A-Z0-9]+)");
    private static final Pattern IDENTIFIER_PROPERTY = Pattern.compile(
            "(?i)^(id|code|identifier|[a-z][a-z0-9_]*(?:_id|_code))$");
    private static final int MAX_MENTIONS = 8;
    private static final int MAX_PROPERTIES = 12;
    private final RagEntityMentionExtractor mentions;
    private final GraphSourceService graphs;

    public GroundingResult ground(Long graphSourceId, String question, GraphSchemaSnapshot authorizedSchema) {
        List<String> identifiers = mentions.extract(question).stream()
                .filter(value -> IDENTIFIER.matcher(value).matches())
                .limit(MAX_MENTIONS).toList();
        if (identifiers.isEmpty()) return GroundingResult.empty();

        List<String> properties = authorizedSchema.nodes().stream()
                .flatMap(node -> node.properties().keySet().stream())
                .filter(name -> IDENTIFIER_PROPERTY.matcher(name).matches())
                .distinct().limit(MAX_PROPERTIES).toList();
        if (properties.isEmpty()) {
            return new GroundingResult(List.of(), List.of("graph schema exposes no identifier-like properties"));
        }

        try {
            Map<String, Object> result = graphs.executeQuery(graphSourceId,
                    lookupCypher(identifiers, properties, authorizedSchema.labelNames()));
            return new GroundingResult(matches(result.get("nodes"), identifiers, properties,
                    new LinkedHashSet<>(authorizedSchema.labelNames())), List.of());
        } catch (Exception error) {
            return new GroundingResult(List.of(),
                    List.of("graph entity grounding failed: " + shortMessage(error)));
        }
    }

    private static String lookupCypher(List<String> identifiers, List<String> properties, List<String> labels) {
        String values = identifiers.stream().map(GraphEntityGroundingService::literal)
                .collect(java.util.stream.Collectors.joining(", "));
        String propertyPredicate = properties.stream()
                .map(property -> "n.`" + property.replace("`", "``") + "` IN [" + values + "]")
                .collect(java.util.stream.Collectors.joining(" OR "));
        String allowedLabels = labels.stream().map(GraphEntityGroundingService::literal)
                .collect(java.util.stream.Collectors.joining(", "));
        return "MATCH (n) WHERE (" + propertyPredicate + ") "
                + "AND any(label IN labels(n) WHERE label IN [" + allowedLabels + "]) "
                + "RETURN n LIMIT 20";
    }

    private static List<EntityMatch> matches(Object rawNodes, List<String> identifiers,
                                             List<String> properties, Set<String> allowedLabels) {
        if (!(rawNodes instanceof Collection<?> nodes)) return List.of();
        List<EntityMatch> result = new ArrayList<>();
        for (Object item : nodes) {
            if (!(item instanceof Map<?, ?> node)) continue;
            List<String> labels = strings(node.get("labels")).stream().filter(allowedLabels::contains).toList();
            if (labels.isEmpty() || !(node.get("properties") instanceof Map<?, ?> values)) continue;
            for (String mention : identifiers) {
                for (String property : properties) {
                    Object value = values.get(property);
                    if (value != null && mention.equalsIgnoreCase(String.valueOf(value))) {
                        result.add(new EntityMatch(mention, labels, property));
                        break;
                    }
                }
            }
        }
        return result.stream().distinct().toList();
    }

    private static List<String> strings(Object value) {
        if (!(value instanceof Collection<?> values)) return List.of();
        return values.stream().map(String::valueOf).toList();
    }

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

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

    public record EntityMatch(String mention, List<String> labels, String property) {
        public EntityMatch { labels = List.copyOf(labels == null ? List.of() : labels); }
        public Map<String, Object> toMetadata() {
            Map<String, Object> result = new LinkedHashMap<>();
            result.put("mention", mention);
            result.put("labels", labels);
            result.put("property", property);
            return result;
        }
        public String toPromptText() {
            return mention + " -> labels " + String.join(",", labels) + " via property " + property;
        }
    }

    public record GroundingResult(List<EntityMatch> matches, List<String> warnings) {
        public GroundingResult {
            matches = List.copyOf(matches == null ? List.of() : matches);
            warnings = List.copyOf(warnings == null ? List.of() : warnings);
        }
        public static GroundingResult empty() { return new GroundingResult(List.of(), List.of()); }
        public List<String> labels() {
            LinkedHashSet<String> result = new LinkedHashSet<>();
            matches.forEach(match -> result.addAll(match.labels()));
            return new ArrayList<>(result);
        }
        public List<Map<String, Object>> metadata() {
            return matches.stream().map(EntityMatch::toMetadata).toList();
        }
        public List<String> promptMentions() {
            return matches.stream().map(EntityMatch::toPromptText).toList();
        }
    }
}
