package at.procon.dip.embedding.service;

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

import at.procon.dip.domain.document.DocumentType;
import at.procon.dip.domain.document.EmbeddingStatus;
import at.procon.dip.domain.document.RepresentationType;
import at.procon.dip.domain.document.entity.Document;
import at.procon.dip.domain.document.entity.DocumentTextRepresentation;
import at.procon.dip.domain.document.repository.DocumentTextRepresentationRepository;
import at.procon.dip.embedding.job.service.EmbeddingJobService;
import at.procon.dip.embedding.model.EmbeddingJobType;
import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
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.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.data.domain.Pageable;

@ExtendWith(MockitoExtension.class)
class ScopedEmbeddingEnqueueServiceTest {

    @Mock
    private DocumentTextRepresentationRepository representationRepository;
    @Mock
    private EmbeddingJobService jobService;
    @Mock
    private EmbeddingModelRegistry modelRegistry;

    private ScopedEmbeddingEnqueueService service;

    @BeforeEach
    void setUp() {
        service = new ScopedEmbeddingEnqueueService(representationRepository, jobService, modelRegistry);
    }

    @Test
    void queuesOnlyMatchingRepresentationsForDocumentType() {
        String modelKey = "e5-default";
        UUID documentId = UUID.randomUUID();
        UUID representationId = UUID.randomUUID();
        Document document = new Document();
        document.setId(documentId);
        DocumentTextRepresentation representation = new DocumentTextRepresentation();
        representation.setId(representationId);
        representation.setDocument(document);

        when(modelRegistry.getRequiredDefaultDocumentModelKey()).thenReturn(modelKey);
        when(modelRegistry.getRequired(modelKey)).thenReturn(anyModel(modelKey));
        when(representationRepository.findEmbeddingCandidatesByDocumentType(
                any(), any(), any(), any(Boolean.class), any(), any(), any(Boolean.class), any(Pageable.class)))
                .thenReturn(List.of(representation));

        ScopedEmbeddingEnqueueResult result = service.enqueueByDocumentType(
                DocumentType.TED_NOTICE_LOT,
                RepresentationType.SEMANTIC_TEXT,
                null,
                true,
                null,
                false,
                250
        );

        verify(representationRepository).findEmbeddingCandidatesByDocumentType(
                DocumentType.TED_NOTICE_LOT,
                RepresentationType.SEMANTIC_TEXT,
                null,
                true,
                modelKey,
                EmbeddingStatus.COMPLETED,
                false,
                Pageable.ofSize(250)
        );
        verify(jobService).enqueueForRepresentation(
                documentId,
                representationId,
                modelKey,
                EmbeddingJobType.DOCUMENT_EMBED
        );
        assertThat(result.jobsQueuedOrAlreadyActive()).isEqualTo(1);
        assertThat(result.matchedRepresentations()).isEqualTo(1);
        assertThat(result.modelKey()).isEqualTo(modelKey);
    }

    @Test
    void capsLargeLimits() {
        String modelKey = "e5-default";
        when(modelRegistry.getRequired(modelKey)).thenReturn(anyModel(modelKey));
        when(representationRepository.findEmbeddingCandidatesByDocumentType(
                any(), any(), any(), any(Boolean.class), any(), any(), any(Boolean.class), any(Pageable.class)))
                .thenReturn(List.of());

        service.enqueueByDocumentType(DocumentType.TED_NOTICE_LOT, null, null, false, modelKey, true, 50_000);

        ArgumentCaptor<Pageable> pageable = ArgumentCaptor.forClass(Pageable.class);
        verify(representationRepository).findEmbeddingCandidatesByDocumentType(
                any(), any(), any(), any(Boolean.class), any(), any(), any(Boolean.class), pageable.capture());
        assertThat(pageable.getValue().getPageSize()).isEqualTo(10_000);
    }

    private EmbeddingModelDescriptor anyModel(String modelKey) {
        return new EmbeddingModelDescriptor(modelKey, "mock", modelKey, 3, null, true, false, null, true, null, null, null);
    }
}
