package com.agent.management.rag.kb;

import com.agent.management.model.entity.RagKnowledgeSourceBinding;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.repository.RagKnowledgeSourceBindingRepository;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.Map;
import java.util.Optional;

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

class RagKnowledgeBaseConfigServiceTest {
    @Test
    void automaticGraphGenerationRejectsEmptyAuthorization() {
        RagKnowledgeSourceBindingRepository repository = mock(RagKnowledgeSourceBindingRepository.class);
        RagKnowledgeSourceBinding binding = binding(RagSourceType.GRAPH);
        when(repository.findById(2L)).thenReturn(Optional.of(binding));
        RagKnowledgeBaseConfigService service = new RagKnowledgeBaseConfigService(repository, new ObjectMapper());

        assertThatThrownBy(() -> service.updateConfig(1L, 2L,
                Map.of("retrievalMode", "AUTO_GENERATE", "allowedLabels", List.of())))
                .isInstanceOf(IllegalArgumentException.class).hasMessageContaining("allowedLabels");
        verify(repository, never()).save(any());
    }

    @Test
    void explicitAuthorizationCanBeSaved() {
        RagKnowledgeSourceBindingRepository repository = mock(RagKnowledgeSourceBindingRepository.class);
        RagKnowledgeSourceBinding binding = binding(RagSourceType.STRUCTURED_DATA);
        when(repository.findById(2L)).thenReturn(Optional.of(binding));
        when(repository.save(binding)).thenReturn(binding);
        RagKnowledgeBaseConfigService service = new RagKnowledgeBaseConfigService(repository, new ObjectMapper());

        service.updateConfig(1L, 2L, Map.of("retrievalMode", "AUTO_GENERATE", "allowedTables", List.of("rescue_forces")));

        verify(repository).save(binding);
    }

    private static RagKnowledgeSourceBinding binding(RagSourceType type) {
        RagKnowledgeSourceBinding binding = new RagKnowledgeSourceBinding();
        binding.setKnowledgeBaseId(1L);
        binding.setSourceType(type);
        return binding;
    }
}
