package at.procon.dip.embedding.policy;

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

import at.procon.dip.domain.document.RepresentationType;
import at.procon.dip.domain.document.entity.DocumentTextRepresentation;
import at.procon.dip.domain.document.repository.DocumentTextRepresentationRepository;
import at.procon.dip.embedding.config.EmbeddingProperties;
import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import at.procon.dip.domain.document.DistanceMetric;
import java.util.List;
import java.util.UUID;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;

@ExtendWith(MockitoExtension.class)
class DefaultEmbeddingSelectionPolicyTest {

    @Mock
    private DocumentTextRepresentationRepository representationRepository;

    private DefaultEmbeddingSelectionPolicy policy;
    private EmbeddingModelDescriptor model;

    @BeforeEach
    void setUp() {
        EmbeddingProperties properties = new EmbeddingProperties();
        properties.getIndexing().setEmbedSemanticText(true);
        properties.getIndexing().setEmbedTitleAbstract(true);
        properties.getIndexing().setEmbedChunks(true);
        properties.getIndexing().setChunkMinLength(20);
        policy = new DefaultEmbeddingSelectionPolicy(representationRepository, properties);
        model = new EmbeddingModelDescriptor("mock-search", "mock-provider", "mock-search", 16,
                DistanceMetric.COSINE, true, false, 4096, true);
    }

    @Test
    void should_select_semantic_text_title_and_large_chunks() {
        UUID documentId = UUID.randomUUID();
        DocumentTextRepresentation semantic = representation(RepresentationType.SEMANTIC_TEXT, true, null,
                "semantic text");
        DocumentTextRepresentation titleAbstract = representation(RepresentationType.TITLE_ABSTRACT, false, null,
                "title abstract text");
        DocumentTextRepresentation shortChunk = representation(RepresentationType.CHUNK, false, 0,
                "too short");
        DocumentTextRepresentation longChunk = representation(RepresentationType.CHUNK, false, 1,
                "This is a sufficiently long chunk that should be selected.");
        when(representationRepository.findByDocument_Id(documentId))
                .thenReturn(List.of(longChunk, semantic, shortChunk, titleAbstract));

        List<DocumentTextRepresentation> selected = policy.selectRepresentations(documentId, model);

        assertThat(selected)
                .extracting(DocumentTextRepresentation::getRepresentationType)
                .containsExactly(RepresentationType.SEMANTIC_TEXT, RepresentationType.TITLE_ABSTRACT, RepresentationType.CHUNK);
        assertThat(selected)
                .extracting(DocumentTextRepresentation::getChunkIndex)
                .containsExactly(null, null, 1);
    }

    private DocumentTextRepresentation representation(RepresentationType type,
                                                      boolean primary,
                                                      Integer chunkIndex,
                                                      String text) {
        return DocumentTextRepresentation.builder()
                .id(UUID.randomUUID())
                .representationType(type)
                .primaryRepresentation(primary)
                .chunkIndex(chunkIndex)
                .charCount(text.length())
                .textBody(text)
                .build();
    }
}
