package at.procon.dip.embedding.service;

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

import at.procon.dip.domain.document.DistanceMetric;
import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import at.procon.dip.embedding.model.EmbeddingProviderResult;
import at.procon.dip.embedding.model.EmbeddingUseCase;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
import java.util.List;
import org.junit.jupiter.api.Test;

class DefaultQueryEmbeddingServiceTest {

    @Test
    void should_use_default_query_model() {
        EmbeddingExecutionService executionService = mock(EmbeddingExecutionService.class);
        EmbeddingModelRegistry modelRegistry = mock(EmbeddingModelRegistry.class);

        when(modelRegistry.getRequiredDefaultQueryModelKey()).thenReturn("mock-search");
        when(executionService.embedTexts("mock-search", EmbeddingUseCase.QUERY, List.of("framework agreement")))
                .thenReturn(new EmbeddingProviderResult(
                        new EmbeddingModelDescriptor("mock-search", "mock-default", "mock-search", 4,
                                DistanceMetric.COSINE, true, true, null, true),
                        List.of(new float[]{1f, 2f, 3f, 4f}),
                        List.of(),
                        "req-1",
                        2
                ));

        DefaultQueryEmbeddingService service = new DefaultQueryEmbeddingService(executionService, modelRegistry);
        float[] vector = service.embedQuery("framework agreement");

        assertThat(vector).containsExactly(1f, 2f, 3f, 4f);
    }
}
