package com.agent.management.rag.memory;

import com.agent.management.model.entity.RagQueryExample;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.repository.RagQueryExampleRepository;
import org.junit.jupiter.api.Test;

import java.util.List;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.*;

class RagExampleMemoryServiceTest {
    @Test
    void retrievesVerifiedExamplesByQuestionSimilarity() {
        RagQueryExampleRepository repository = mock(RagQueryExampleRepository.class);
        when(repository.findBySourceTypeAndSourceIdAndVerifiedTrue(RagSourceType.STRUCTURED_DATA, "1"))
                .thenReturn(List.of(
                        example("查询可用救援力量", "SELECT * FROM rescue_forces WHERE status='READY'"),
                        example("统计审计日志", "SELECT COUNT(*) FROM audit_log")
                ));
        var profiles=mock(com.agent.management.rag.capability.RagCapabilityProfileService.class); when(profiles.get(any(),anyString())).thenReturn(new com.agent.management.rag.capability.RagCapabilityProfile(RagSourceType.STRUCTURED_DATA,"1","v",java.time.Instant.now(),null,null,List.of())); RagExampleMemoryService service = new RagExampleMemoryService(repository, mock(com.agent.management.rag.capability.RagSemanticCatalogService.class),profiles);

        assertThat(service.findSimilar(RagSourceType.STRUCTURED_DATA, "1", "当前有哪些可用救援力量", 1))
                .singleElement().satisfies(item -> assertThat(item.get("query")).isEqualTo("SELECT * FROM rescue_forces WHERE status='READY'"));
    }

    @Test
    void learnedExampleRequiresReviewBeforeReuse() {
        RagQueryExampleRepository repository = mock(RagQueryExampleRepository.class);
        when(repository.findFirstBySourceTypeAndSourceIdAndQuestionAndQueryText(any(), anyString(), anyString(), anyString()))
                .thenReturn(java.util.Optional.empty());
        when(repository.save(any())).thenAnswer(invocation -> invocation.getArgument(0));
        var profiles=mock(com.agent.management.rag.capability.RagCapabilityProfileService.class); when(profiles.get(any(),anyString())).thenReturn(new com.agent.management.rag.capability.RagCapabilityProfile(RagSourceType.GRAPH,"1","v",java.time.Instant.now(),null,null,List.of())); RagExampleMemoryService service = new RagExampleMemoryService(repository, mock(com.agent.management.rag.capability.RagSemanticCatalogService.class),profiles);

        service.recordSuccess(RagSourceType.GRAPH, "1", "任务阶段", "MATCH (n) RETURN n");

        org.mockito.ArgumentCaptor<RagQueryExample> captor = org.mockito.ArgumentCaptor.forClass(RagQueryExample.class);
        verify(repository).save(captor.capture());
        assertThat(captor.getValue().isVerified()).isFalse();
    }

    @Test
    void corruptedQuestionIsNotLearned() {
        RagQueryExampleRepository repository = mock(RagQueryExampleRepository.class);
        RagExampleMemoryService service = new RagExampleMemoryService(repository, mock(com.agent.management.rag.capability.RagSemanticCatalogService.class),mock(com.agent.management.rag.capability.RagCapabilityProfileService.class));

        service.recordSuccess(RagSourceType.GRAPH, "1", "????????????", "MATCH (n) RETURN n");

        verifyNoInteractions(repository);
    }

    private static RagQueryExample example(String question, String query) {
        RagQueryExample example = new RagQueryExample();
        example.setQuestion(question); example.setQueryText(query); example.setVerified(true);
        return example;
    }
}
