package com.agent.management.rag.capability;

import com.agent.management.model.entity.RagGraphGovernanceConfigEntity;
import com.agent.management.rag.model.RagQuery;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.repository.RagGraphGovernanceConfigRepository;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.time.Instant;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

@Service
@RequiredArgsConstructor
public class GraphGovernanceConfigService {
    private final RagGraphGovernanceConfigRepository repository;
    private final RagCapabilityProfileService profiles;
    private final ObjectMapper objectMapper;

    public Map<String, Object> get(Long graphSourceId) {
        return repository.findByGraphSourceId(graphSourceId).map(this::response)
                .orElseGet(() -> Map.of("graphSourceId", graphSourceId, "configured", false,
                        "config", Map.of()));
    }

    @Transactional
    public Map<String, Object> save(Long graphSourceId, Map<String, Object> config) {
        RagCapabilityProfile profile = profiles.get(RagSourceType.GRAPH, String.valueOf(graphSourceId));
        Map<String, Object> normalized = normalize(config, profile);
        RagGraphGovernanceConfigEntity entity = repository.findByGraphSourceId(graphSourceId)
                .orElseGet(RagGraphGovernanceConfigEntity::new);
        entity.setGraphSourceId(graphSourceId);
        entity.setSchemaVersion(profile.version());
        entity.setUpdatedAt(Instant.now());
        try {
            entity.setConfigJson(objectMapper.writeValueAsString(normalized));
        } catch (Exception error) {
            throw new IllegalArgumentException("invalid graph governance config: " + error.getMessage(), error);
        }
        return response(repository.save(entity));
    }

    public RagQuery apply(RagQuery query, Long graphSourceId) {
        Map<String, Object> stored = config(graphSourceId);
        if (stored.isEmpty()) return query;
        RagQuery effective = new RagQuery();
        effective.setQuery(query.getQuery());
        effective.setUserId(query.getUserId());
        effective.setSessionId(query.getSessionId());
        effective.setSourceIds(query.getSourceIds());
        effective.setTopK(query.getTopK());
        Map<String, Object> filters = new LinkedHashMap<>(stored);
        if (query.getFilters() != null) filters.putAll(query.getFilters());
        effective.setFilters(filters);
        return effective;
    }

    private Map<String, Object> config(Long graphSourceId) {
        return repository.findByGraphSourceId(graphSourceId).map(this::parseConfig).orElseGet(Map::of);
    }

    private static Map<String, Object> normalize(Map<String, Object> raw, RagCapabilityProfile profile) {
        Map<String, Object> source = raw == null ? Map.of() : raw;
        List<String> labels = strings(source.get("allowedLabels")).stream()
                .filter(profile.graphSchema().labelNames()::contains).toList();
        if (labels.isEmpty()) throw new IllegalArgumentException("at least one real graph label must be authorized");
        List<String> relationships = strings(source.get("allowedRelationships")).stream()
                .filter(profile.graphSchema().relationshipTypeNames()::contains).toList();
        Map<String, List<String>> actualProperties = new LinkedHashMap<>();
        profile.graphSchema().nodes().forEach(node -> actualProperties.put(node.label(), List.copyOf(node.properties().keySet())));
        profile.graphSchema().relationships().forEach(rel -> actualProperties.put(rel.type(), List.copyOf(rel.properties().keySet())));
        Map<String, List<String>> properties = new LinkedHashMap<>();
        if (source.get("allowedProperties") instanceof Map<?, ?> values) {
            values.forEach((owner, rawValues) -> {
                String name = String.valueOf(owner);
                if ((!labels.contains(name) && !relationships.contains(name)) || !actualProperties.containsKey(name)) return;
                properties.put(name, strings(rawValues).stream().filter(actualProperties.get(name)::contains).toList());
            });
        }
        Map<String, Object> result = new LinkedHashMap<>();
        result.put("retrievalMode", "AUTO_GENERATE");
        result.put("allowTextToCypher", true);
        result.put("allowedLabels", labels);
        result.put("allowedRelationships", relationships);
        result.put("allowedProperties", properties);
        result.put("generationRules", strings(source.get("generationRules")));
        result.put("intentKeywords", stringMap(source.get("intentKeywords")));
        result.put("genericTerms", strings(source.get("genericTerms")));
        result.put("enabledGenericQueryModes", strings(source.get("enabledGenericQueryModes")));
        result.put("maxDepth", source.get("maxDepth") instanceof Number number
                ? Math.max(1, Math.min(5, number.intValue())) : 3);
        return result;
    }

    private Map<String, Object> response(RagGraphGovernanceConfigEntity entity) {
        return Map.of("graphSourceId", entity.getGraphSourceId(), "configured", true,
                "schemaVersion", entity.getSchemaVersion(), "updatedAt", entity.getUpdatedAt(),
                "config", parseConfig(entity));
    }

    private Map<String, Object> parseConfig(RagGraphGovernanceConfigEntity entity) {
        try {
            return objectMapper.readValue(entity.getConfigJson(), new TypeReference<Map<String, Object>>() {});
        } catch (Exception ignored) {
            return Map.of();
        }
    }

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

    private static Map<String, List<String>> stringMap(Object value) {
        if (!(value instanceof Map<?, ?> map)) return Map.of();
        Map<String, List<String>> result = new LinkedHashMap<>();
        map.forEach((key, raw) -> result.put(String.valueOf(key), strings(raw)));
        return result;
    }
}
