package com.agent.management.rag.capability;

import com.agent.management.rag.graph.GraphSchemaDiscoveryService;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.service.DataSourceService;
import com.agent.management.service.SchemaExplorerService;
import com.agent.management.service.impl.EmbeddingBridgeClient;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.time.Instant;
import java.util.AbstractMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
@RequiredArgsConstructor
public class RagCapabilityProfileService {
    private static final Duration TTL = Duration.ofMinutes(5);
    private final SchemaExplorerService sqlExplorer;
    private final DataSourceService dataSources;
    private final GraphSchemaDiscoveryService graphSchemas;
    private final EmbeddingBridgeClient embeddings;
    private final Map<String, RagCapabilityProfile> cache = new ConcurrentHashMap<>();

    public RagCapabilityProfile get(RagSourceType type, String sourceId) {
        String key = type + ":" + sourceId;
        RagCapabilityProfile cached = cache.get(key);
        if (cached != null && cached.generatedAt().plus(TTL).isAfter(Instant.now())) return cached;
        return refresh(type, sourceId);
    }

    public RagCapabilityProfile refresh(RagSourceType type, String sourceId) {
        Long id = Long.valueOf(sourceId);
        RagCapabilityProfile profile;
        if (type == RagSourceType.STRUCTURED_DATA) {
            var ds = dataSources.getDecrypted(id);
            List<SqlSchemaSnapshot.TableSchema> tables = sqlExplorer.listTables(id, ds.getSchemaName()).stream()
                    .map(table -> new SqlSchemaSnapshot.TableSchema(table.name(), table.remarks(),
                            sqlExplorer.describeTable(id, ds.getSchemaName(), table.name()).stream()
                                    .map(column -> new SqlSchemaSnapshot.ColumnSchema(column.name(), column.type(), column.remarks()))
                                    .toList()))
                    .toList();
            List<SqlSchemaSnapshot.ForeignKey> foreignKeys = tables.stream().flatMap(table -> sqlExplorer.listForeignKeys(id, ds.getSchemaName(), table.name()).stream())
                    .map(key -> new SqlSchemaSnapshot.ForeignKey(key.name(),key.sourceTable(),key.sourceColumn(),key.targetTable(),key.targetColumn())).toList();
            SqlSchemaSnapshot sql = new SqlSchemaSnapshot(tables, foreignKeys);
            List<RagSemanticEntry> catalog = vectorize(tables.stream().map(table ->
                    new AbstractMap.SimpleEntry<>("table:" + table.name(), table.name() + " " + table.description() + " " +
                            table.columns().stream().map(column -> column.name()+" "+column.description()).collect(java.util.stream.Collectors.joining(" ")))).toList());
            profile = new RagCapabilityProfile(type, sourceId, Integer.toHexString(sql.hashCode()), Instant.now(), sql, null, catalog);
        } else if (type == RagSourceType.GRAPH) {
            var graph = graphSchemas.refresh(id);
            List<Map.Entry<String,String>> texts = new java.util.ArrayList<>();
            graph.nodes().forEach(node -> texts.add(Map.entry("label:"+node.label(), node.label()+" "+String.join(" ",node.properties().keySet()))));
            graph.relationships().forEach(rel -> texts.add(Map.entry("relationship:"+rel.type(), rel.type()+" "+String.join(" ",rel.startLabels())+" "+String.join(" ",rel.endLabels()))));
            profile = new RagCapabilityProfile(type, sourceId, Integer.toHexString(graph.hashCode()), Instant.now(), null, graph, vectorize(texts));
        } else {
            throw new IllegalArgumentException("capability profile is unsupported for source type: " + type);
        }
        cache.put(type + ":" + sourceId, profile);
        return profile;
    }

    private List<RagSemanticEntry> vectorize(List<? extends Map.Entry<String,String>> items) {
        if (items.isEmpty()) return List.of();
        try {
            List<List<Double>> vectors = embeddings.embed(items.stream().map(Map.Entry::getValue).toList());
            List<RagSemanticEntry> result = new java.util.ArrayList<>();
            for (int i=0;i<items.size();i++) result.add(new RagSemanticEntry(items.get(i).getKey(),items.get(i).getValue(),vectors.get(i)));
            return result;
        } catch (Exception ignored) {
            return items.stream().map(item -> new RagSemanticEntry(item.getKey(),item.getValue(),List.of())).toList();
        }
    }
}
