package com.agent.management.rag.capability;

import com.agent.management.rag.graph.GraphSchemaSnapshot;
import org.junit.jupiter.api.Test;

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

import static org.assertj.core.api.Assertions.assertThat;

class RagSchemaLinkerTest {
    private final RagSchemaLinker linker = new RagSchemaLinker();

    @Test
    void selectsRelevantSqlTableFromQuestionAndSemanticHints() {
        SqlSchemaSnapshot schema = new SqlSchemaSnapshot(List.of(
                new SqlSchemaSnapshot.TableSchema("rescue_forces", "救援力量和可调度队伍", List.of(
                        new SqlSchemaSnapshot.ColumnSchema("status", "VARCHAR", "调度状态 READY 表示可用"),
                        new SqlSchemaSnapshot.ColumnSchema("capacity", "INT", "载员能力")
                )),
                new SqlSchemaSnapshot.TableSchema("audit_log", "系统审计记录", List.of(
                        new SqlSchemaSnapshot.ColumnSchema("created_at", "TIMESTAMP", "创建时间")
                ))
        ));

        assertThat(linker.selectSqlTables("有哪些可用救援力量并按载员能力排序", schema, List.of(), Map.of()))
                .containsExactly("rescue_forces");
    }

    @Test
    void similarGraphExampleSelectsItsSchemaSubgraph() {
        GraphSchemaSnapshot graph = new GraphSchemaSnapshot(
                List.of(
                        new GraphSchemaSnapshot.NodeSchema("Mission", Map.of("name", List.of("String"))),
                        new GraphSchemaSnapshot.NodeSchema("Stage", Map.of("name", List.of("String"))),
                        new GraphSchemaSnapshot.NodeSchema("Secret", Map.of("value", List.of("String")))
                ),
                List.of(
                        new GraphSchemaSnapshot.RelationshipSchema("HAS_STAGE", List.of("Mission"), List.of("Stage"), Map.of()),
                        new GraphSchemaSnapshot.RelationshipSchema("SECRET_LINK", List.of("Mission"), List.of("Secret"), Map.of())
                )
        );
        List<Map<String, Object>> examples = List.of(Map.of(
                "question", "任务包括哪些阶段",
                "cypher", "MATCH (:Mission)-[:HAS_STAGE]->(:Stage) RETURN 1"
        ));

        GraphSchemaSnapshot selected = linker.selectGraphSchema("东海任务阶段是什么", graph, examples);

        assertThat(selected.labelNames()).containsExactlyInAnyOrder("Mission", "Stage");
        assertThat(selected.relationshipTypeNames()).containsExactly("HAS_STAGE");
    }
}
