package com.agent.management.rag.kb;

import com.agent.management.model.entity.RagKnowledgeBase;
import com.agent.management.model.entity.RagKnowledgeSourceBinding;
import com.agent.management.rag.api.RagRetriever;
import com.agent.management.rag.model.*;
import com.agent.management.repository.RagKnowledgeBaseRepository;
import com.agent.management.repository.RagKnowledgeSourceBindingRepository;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

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

class KnowledgeBaseRagRetrieverTest {
    @Test
    void bindingResolvesSourceAndScopedFiltersOverrideConfigWithoutDefaults() {
        RagKnowledgeBaseRepository bases = mock(RagKnowledgeBaseRepository.class);
        RagKnowledgeSourceBindingRepository bindings = mock(RagKnowledgeSourceBindingRepository.class);
        RagRetriever structured = mock(RagRetriever.class);
        when(structured.sourceType()).thenReturn(RagSourceType.STRUCTURED_DATA);
        RagRetrievalResult rr = new RagRetrievalResult();
        rr.setSourceType(RagSourceType.STRUCTURED_DATA);
        when(structured.retrieve(any())).thenReturn(rr);

        RagKnowledgeBase kb = new RagKnowledgeBase();
        kb.setId(1L);
        kb.setEnabled(true);
        RagKnowledgeSourceBinding binding = new RagKnowledgeSourceBinding();
        binding.setSourceType(RagSourceType.STRUCTURED_DATA);
        binding.setSourceId("9");
        binding.setEnabled(true);
        binding.setConfigJson("{\"allowTextToSql\":true,\"maxRows\":50,\"tableWhitelist\":[\"rescue_forces\"]}");
        when(bases.findById(1L)).thenReturn(Optional.of(kb));
        when(bindings.findByKnowledgeBaseIdAndEnabledTrueOrderByPriorityAsc(1L)).thenReturn(List.of(binding));

        KnowledgeBaseRagRequest request = new KnowledgeBaseRagRequest();
        request.setKnowledgeBaseId(1L);
        request.setQuery("q");
        request.setFilters(Map.of("structured", Map.of("maxRows", 10)));
        RagQuestionPlanningService planner=mock(RagQuestionPlanningService.class);when(planner.plan(anyString(),anyList())).thenReturn(Map.of());
        new KnowledgeBaseRagRetriever(bases, bindings, List.of(structured), new ObjectMapper(), Runnable::run,planner).retrieve(request);

        ArgumentCaptor<RagQuery> captor = ArgumentCaptor.forClass(RagQuery.class);
        verify(structured).retrieve(captor.capture());
        assertThat(captor.getValue().getSourceIds()).containsExactly("9");
        assertThat(captor.getValue().getFilters())
                .containsEntry("allowTextToSql", true)
                .containsEntry("maxRows", 10)
                .doesNotContainKeys("defaultSql", "sql");
    }

    @Test
    void retrievesMultipleBindingsConcurrently() throws Exception {
        RagKnowledgeBaseRepository bases = mock(RagKnowledgeBaseRepository.class);
        RagKnowledgeSourceBindingRepository bindings = mock(RagKnowledgeSourceBindingRepository.class);
        RagKnowledgeBase kb = new RagKnowledgeBase(); kb.setId(1L); kb.setEnabled(true);
        when(bases.findById(1L)).thenReturn(Optional.of(kb));
        RagKnowledgeSourceBinding structuredBinding = binding(RagSourceType.STRUCTURED_DATA, "1", 10);
        RagKnowledgeSourceBinding graphBinding = binding(RagSourceType.GRAPH, "2", 20);
        when(bindings.findByKnowledgeBaseIdAndEnabledTrueOrderByPriorityAsc(1L))
                .thenReturn(List.of(structuredBinding, graphBinding));

        CountDownLatch bothStarted = new CountDownLatch(2);
        CountDownLatch release = new CountDownLatch(1);
        RagRetriever structured = blockingRetriever(RagSourceType.STRUCTURED_DATA, bothStarted, release);
        RagRetriever graph = blockingRetriever(RagSourceType.GRAPH, bothStarted, release);
        KnowledgeBaseRagRequest request = new KnowledgeBaseRagRequest();
        request.setKnowledgeBaseId(1L); request.setQuery("q");
        RagQuestionPlanningService concurrentPlanner=mock(RagQuestionPlanningService.class);when(concurrentPlanner.plan(anyString(),anyList())).thenReturn(Map.of());
        KnowledgeBaseRagRetriever subject = new KnowledgeBaseRagRetriever(
                bases, bindings, List.of(structured, graph), new ObjectMapper(),
                command -> new Thread(command, "test-rag-retrieval").start(), concurrentPlanner);

        Thread invocation = new Thread(() -> subject.retrieve(request));
        invocation.start();
        assertThat(bothStarted.await(1, TimeUnit.SECONDS)).isTrue();
        release.countDown();
        invocation.join(2000);
        assertThat(invocation.isAlive()).isFalse();
    }

    private static RagKnowledgeSourceBinding binding(RagSourceType type, String sourceId, int priority) {
        RagKnowledgeSourceBinding binding = new RagKnowledgeSourceBinding();
        binding.setSourceType(type); binding.setSourceId(sourceId); binding.setPriority(priority); binding.setEnabled(true);
        return binding;
    }

    private static RagRetriever blockingRetriever(RagSourceType type, CountDownLatch started, CountDownLatch release) {
        return new RagRetriever() {
            @Override public RagSourceType sourceType() { return type; }
            @Override public RagRetrievalResult retrieve(RagQuery query) {
                started.countDown();
                try { release.await(2, TimeUnit.SECONDS); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }
                RagRetrievalResult result = new RagRetrievalResult(); result.setSourceType(type); return result;
            }
        };
    }
}
