package at.procon.dip.embedding.service;

import at.procon.dip.domain.document.DistanceMetric;
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.config.EmbeddingProperties;
import at.procon.dip.embedding.job.entity.EmbeddingJob;
import at.procon.dip.embedding.job.service.EmbeddingJobService;
import at.procon.dip.embedding.model.*;
import at.procon.dip.embedding.policy.EmbeddingSelectionPolicy;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
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 java.util.List;
import java.util.Optional;
import java.util.UUID;

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

@ExtendWith(MockitoExtension.class)
class RepresentationEmbeddingOrchestratorTest {

    @Mock
    private EmbeddingJobService jobService;
    @Mock
    private EmbeddingExecutionService executionService;
    @Mock
    private EmbeddingJobExecutionPersistenceService executionPersistenceService;
    @Mock
    private DocumentTextRepresentationRepository representationRepository;
    @Mock
    private EmbeddingSelectionPolicy selectionPolicy;
    @Mock
    private EmbeddingModelRegistry modelRegistry;

    private EmbeddingProperties properties;
    private RepresentationEmbeddingOrchestrator orchestrator;

    @BeforeEach
    void setUp() {
        properties = new EmbeddingProperties();
        properties.setEnabled(true);
        properties.getJobs().setEnabled(true);
        properties.getJobs().setBatchSize(10);
        properties.getJobs().setExecutionBatchSize(10);
        properties.getIndexing().setFallbackMaxInputChars(8192);
        orchestrator = new RepresentationEmbeddingOrchestrator(
                jobService,
                executionService,
                executionPersistenceService,
                representationRepository,
                selectionPolicy,
                modelRegistry,
                properties,
                Runnable::run
        );
    }

    @Test
    void processClaimedJob_should_embed_representation_and_mark_job_completed() {
        UUID documentId = UUID.randomUUID();
        UUID representationId = UUID.randomUUID();
        UUID embeddingId = UUID.randomUUID();
        EmbeddingJob job = EmbeddingJob.builder()
                .id(UUID.randomUUID())
                .documentId(documentId)
                .representationId(representationId)
                .modelKey("mock-search")
                .jobType(EmbeddingJobType.DOCUMENT_EMBED)
                .status(EmbeddingJobStatus.IN_PROGRESS)
                .attemptCount(1)
                .build();

        Document document = Document.builder().id(documentId).build();
        DocumentTextRepresentation representation = DocumentTextRepresentation.builder()
                .id(representationId)
                .document(document)
                .textBody("District heating optimization strategy")
                .build();
        EmbeddingModelDescriptor model = new EmbeddingModelDescriptor(
                "mock-search", "mock-provider", "mock-search", 16,
                DistanceMetric.COSINE, true, false, 4096, true
        );
        when(representationRepository.findById(representationId)).thenReturn(Optional.of(representation));
        when(modelRegistry.getRequired("mock-search")).thenReturn(model);
        when(executionPersistenceService.startProcessing(representationId, "mock-search"))
                .thenReturn(embeddingId);
        when(executionService.embedTexts(eq("mock-search"), eq(EmbeddingUseCase.DOCUMENT), any()))
                .thenReturn(new EmbeddingProviderResult(
                        model,
                        List.of(new float[]{0.1f, 0.2f}),
                        List.of(),
                        "req-1",
                        12
                ));

        orchestrator.processClaimedJob(job);

        verify(executionPersistenceService).startProcessing(representationId, "mock-search");
        verify(executionPersistenceService).completeJob(eq(embeddingId), any(EmbeddingProviderResult.class), eq(job.getId()), eq("req-1"));
    }

    @Test
    void processNextReadyBatch_should_group_batchable_jobs_and_call_provider_once() {
        properties.getJobs().setProcessInBatches(true);
        properties.getJobs().setExecutionBatchSize(10);

        UUID documentId = UUID.randomUUID();
        UUID representationId1 = UUID.randomUUID();
        UUID representationId2 = UUID.randomUUID();
        UUID embeddingId1 = UUID.randomUUID();
        UUID embeddingId2 = UUID.randomUUID();

        EmbeddingJob job1 = EmbeddingJob.builder()
                .id(UUID.randomUUID())
                .documentId(documentId)
                .representationId(representationId1)
                .modelKey("e5-default")
                .jobType(EmbeddingJobType.DOCUMENT_EMBED)
                .status(EmbeddingJobStatus.IN_PROGRESS)
                .attemptCount(1)
                .build();
        EmbeddingJob job2 = EmbeddingJob.builder()
                .id(UUID.randomUUID())
                .documentId(documentId)
                .representationId(representationId2)
                .modelKey("e5-default")
                .jobType(EmbeddingJobType.DOCUMENT_EMBED)
                .status(EmbeddingJobStatus.IN_PROGRESS)
                .attemptCount(1)
                .build();

        when(jobService.claimNextReadyJobs(10)).thenReturn(List.of(job1, job2), List.of());

        Document document = Document.builder().id(documentId).build();
        when(representationRepository.findById(representationId1)).thenReturn(Optional.of(
                DocumentTextRepresentation.builder().id(representationId1).document(document).textBody("alpha").build()
        ));
        when(representationRepository.findById(representationId2)).thenReturn(Optional.of(
                DocumentTextRepresentation.builder().id(representationId2).document(document).textBody("beta").build()
        ));

        EmbeddingModelDescriptor model = new EmbeddingModelDescriptor(
                "e5-default", "vector-sync-e5", "intfloat/multilingual-e5-large", 3,
                DistanceMetric.COSINE, true, true, 4096, true
        );
        when(modelRegistry.getRequired("e5-default")).thenReturn(model);
        when(executionPersistenceService.startProcessing(representationId1, "e5-default"))
                .thenReturn(embeddingId1);
        when(executionPersistenceService.startProcessing(representationId2, "e5-default"))
                .thenReturn(embeddingId2);
        when(executionService.embedTexts(eq("e5-default"), eq(EmbeddingUseCase.DOCUMENT), any()))
                .thenReturn(new EmbeddingProviderResult(
                        model,
                        List.of(new float[]{0.1f, 0.2f, 0.3f}, new float[]{0.4f, 0.5f, 0.6f}),
                        List.of(),
                        "batch-req-1",
                        21
                ));

        int processed = orchestrator.processNextReadyBatch();

        assertThat(processed).isEqualTo(2);
        verify(executionPersistenceService).startProcessing(representationId1, "e5-default");
        verify(executionPersistenceService).startProcessing(representationId2, "e5-default");
        ArgumentCaptor<List<String>> textsCaptor = ArgumentCaptor.forClass(List.class);
        verify(executionService, times(1)).embedTexts(eq("e5-default"), eq(EmbeddingUseCase.DOCUMENT), textsCaptor.capture());
        assertThat(textsCaptor.getValue()).containsExactly("alpha", "beta");
        verify(executionPersistenceService).completeJob(eq(embeddingId1), aryEq(new float[]{0.1f, 0.2f, 0.3f}), eq(null), eq(job1.getId()), eq("batch-req-1"), eq(EmbeddingPrefixMode.OFF), eq(null));
        verify(executionPersistenceService).completeJob(eq(embeddingId2), aryEq(new float[]{0.4f, 0.5f, 0.6f}), eq(null), eq(job2.getId()), eq("batch-req-1"), eq(EmbeddingPrefixMode.OFF), eq(null));
    }

    @Test
    void processNextReadyBatch_should_fall_back_to_single_processing_for_non_batch_model() {
        properties.getJobs().setProcessInBatches(true);

        UUID documentId = UUID.randomUUID();
        UUID representationId = UUID.randomUUID();
        UUID embeddingId = UUID.randomUUID();
        EmbeddingJob job = EmbeddingJob.builder()
                .id(UUID.randomUUID())
                .documentId(documentId)
                .representationId(representationId)
                .modelKey("mock-search")
                .jobType(EmbeddingJobType.DOCUMENT_EMBED)
                .status(EmbeddingJobStatus.IN_PROGRESS)
                .attemptCount(1)
                .build();

        when(jobService.claimNextReadyJobs(10)).thenReturn(List.of(job), List.of());
        Document document = Document.builder().id(documentId).build();
        when(representationRepository.findById(representationId)).thenReturn(Optional.of(
                DocumentTextRepresentation.builder().id(representationId).document(document).textBody("gamma").build()
        ));
        EmbeddingModelDescriptor model = new EmbeddingModelDescriptor(
                "mock-search", "mock-provider", "mock-search", 2,
                DistanceMetric.COSINE, true, false, 4096, true
        );
        when(modelRegistry.getRequired("mock-search")).thenReturn(model);
        when(executionPersistenceService.startProcessing(representationId, "mock-search"))
                .thenReturn(embeddingId);
        when(executionService.embedTexts(eq("mock-search"), eq(EmbeddingUseCase.DOCUMENT), any()))
                .thenReturn(new EmbeddingProviderResult(
                        model,
                        List.of(new float[]{0.7f, 0.8f}),
                        List.of(),
                        "req-2",
                        7
                ));

        orchestrator.processNextReadyBatch();

        verify(executionService, times(1)).embedTexts(eq("mock-search"), eq(EmbeddingUseCase.DOCUMENT), eq(List.of("gamma")));
        verify(executionPersistenceService, never()).completeJob(eq(embeddingId), any(float[].class), eq(null), eq(job.getId()), anyString(), any(), any());
        verify(executionPersistenceService).completeJob(eq(embeddingId), any(EmbeddingProviderResult.class), eq(job.getId()), eq("req-2"));
    }

    @Test
    void processNextReadyBatch_should_keep_claiming_new_batches_until_all_parallel_slots_are_exhausted() {
        properties.getJobs().setParallelBatchCount(2);
        properties.getJobs().setBatchSize(2);
        properties.getJobs().setProcessInBatches(true);
        properties.getJobs().setExecutionBatchSize(2);

        UUID documentId = UUID.randomUUID();
        UUID representationId1 = UUID.randomUUID();
        UUID representationId2 = UUID.randomUUID();
        UUID representationId3 = UUID.randomUUID();
        UUID representationId4 = UUID.randomUUID();
        UUID representationId5 = UUID.randomUUID();
        UUID representationId6 = UUID.randomUUID();
        UUID embeddingId1 = UUID.randomUUID();
        UUID embeddingId2 = UUID.randomUUID();
        UUID embeddingId3 = UUID.randomUUID();
        UUID embeddingId4 = UUID.randomUUID();
        UUID embeddingId5 = UUID.randomUUID();
        UUID embeddingId6 = UUID.randomUUID();

        EmbeddingJob job1 = EmbeddingJob.builder().id(UUID.randomUUID()).documentId(documentId).representationId(representationId1).modelKey("e5-default").jobType(EmbeddingJobType.DOCUMENT_EMBED).status(EmbeddingJobStatus.IN_PROGRESS).attemptCount(1).build();
        EmbeddingJob job2 = EmbeddingJob.builder().id(UUID.randomUUID()).documentId(documentId).representationId(representationId2).modelKey("e5-default").jobType(EmbeddingJobType.DOCUMENT_EMBED).status(EmbeddingJobStatus.IN_PROGRESS).attemptCount(1).build();
        EmbeddingJob job3 = EmbeddingJob.builder().id(UUID.randomUUID()).documentId(documentId).representationId(representationId3).modelKey("e5-default").jobType(EmbeddingJobType.DOCUMENT_EMBED).status(EmbeddingJobStatus.IN_PROGRESS).attemptCount(1).build();
        EmbeddingJob job4 = EmbeddingJob.builder().id(UUID.randomUUID()).documentId(documentId).representationId(representationId4).modelKey("e5-default").jobType(EmbeddingJobType.DOCUMENT_EMBED).status(EmbeddingJobStatus.IN_PROGRESS).attemptCount(1).build();
        EmbeddingJob job5 = EmbeddingJob.builder().id(UUID.randomUUID()).documentId(documentId).representationId(representationId5).modelKey("e5-default").jobType(EmbeddingJobType.DOCUMENT_EMBED).status(EmbeddingJobStatus.IN_PROGRESS).attemptCount(1).build();
        EmbeddingJob job6 = EmbeddingJob.builder().id(UUID.randomUUID()).documentId(documentId).representationId(representationId6).modelKey("e5-default").jobType(EmbeddingJobType.DOCUMENT_EMBED).status(EmbeddingJobStatus.IN_PROGRESS).attemptCount(1).build();

        when(jobService.claimNextReadyJobs(2))
                .thenReturn(List.of(job1, job2))
                .thenReturn(List.of(job3, job4))
                .thenReturn(List.of(job5, job6))
                .thenReturn(List.of());

        Document document = Document.builder().id(documentId).build();
        when(representationRepository.findById(representationId1)).thenReturn(Optional.of(DocumentTextRepresentation.builder().id(representationId1).document(document).textBody("alpha").build()));
        when(representationRepository.findById(representationId2)).thenReturn(Optional.of(DocumentTextRepresentation.builder().id(representationId2).document(document).textBody("beta").build()));
        when(representationRepository.findById(representationId3)).thenReturn(Optional.of(DocumentTextRepresentation.builder().id(representationId3).document(document).textBody("gamma").build()));
        when(representationRepository.findById(representationId4)).thenReturn(Optional.of(DocumentTextRepresentation.builder().id(representationId4).document(document).textBody("delta").build()));
        when(representationRepository.findById(representationId5)).thenReturn(Optional.of(DocumentTextRepresentation.builder().id(representationId5).document(document).textBody("epsilon").build()));
        when(representationRepository.findById(representationId6)).thenReturn(Optional.of(DocumentTextRepresentation.builder().id(representationId6).document(document).textBody("zeta").build()));

        EmbeddingModelDescriptor model = new EmbeddingModelDescriptor(
                "e5-default", "vector-sync-e5", "intfloat/multilingual-e5-large", 3,
                DistanceMetric.COSINE, true, true, 4096, true
        );
        when(modelRegistry.getRequired("e5-default")).thenReturn(model);
        when(executionPersistenceService.startProcessing(representationId1, "e5-default")).thenReturn(embeddingId1);
        when(executionPersistenceService.startProcessing(representationId2, "e5-default")).thenReturn(embeddingId2);
        when(executionPersistenceService.startProcessing(representationId3, "e5-default")).thenReturn(embeddingId3);
        when(executionPersistenceService.startProcessing(representationId4, "e5-default")).thenReturn(embeddingId4);
        when(executionPersistenceService.startProcessing(representationId5, "e5-default")).thenReturn(embeddingId5);
        when(executionPersistenceService.startProcessing(representationId6, "e5-default")).thenReturn(embeddingId6);
        when(executionService.embedTexts(eq("e5-default"), eq(EmbeddingUseCase.DOCUMENT), any()))
                .thenReturn(new EmbeddingProviderResult(
                        model,
                        List.of(new float[]{0.1f, 0.2f, 0.3f}, new float[]{0.4f, 0.5f, 0.6f}),
                        List.of(),
                        "batch-req-1",
                        21
                ))
                .thenReturn(new EmbeddingProviderResult(
                        model,
                        List.of(new float[]{0.7f, 0.8f, 0.9f}, new float[]{1.0f, 1.1f, 1.2f}),
                        List.of(),
                        "batch-req-2",
                        25
                ))
                .thenReturn(new EmbeddingProviderResult(
                        model,
                        List.of(new float[]{1.3f, 1.4f, 1.5f}, new float[]{1.6f, 1.7f, 1.8f}),
                        List.of(),
                        "batch-req-3",
                        29
                ));

        int processed = orchestrator.processNextReadyBatch();

        assertThat(processed).isEqualTo(6);
        verify(jobService, times(4)).claimNextReadyJobs(2);
        verify(executionService, times(3)).embedTexts(eq("e5-default"), eq(EmbeddingUseCase.DOCUMENT), any());
        verify(executionPersistenceService).completeJob(eq(embeddingId1), aryEq(new float[]{0.1f, 0.2f, 0.3f}), eq(null), eq(job1.getId()), eq("batch-req-1"), eq(EmbeddingPrefixMode.OFF), eq(null));
        verify(executionPersistenceService).completeJob(eq(embeddingId2), aryEq(new float[]{0.4f, 0.5f, 0.6f}), eq(null), eq(job2.getId()), eq("batch-req-1"), eq(EmbeddingPrefixMode.OFF), eq(null));
        verify(executionPersistenceService).completeJob(eq(embeddingId3), aryEq(new float[]{0.7f, 0.8f, 0.9f}), eq(null), eq(job3.getId()), eq("batch-req-2"), eq(EmbeddingPrefixMode.OFF), eq(null));
        verify(executionPersistenceService).completeJob(eq(embeddingId4), aryEq(new float[]{1.0f, 1.1f, 1.2f}), eq(null), eq(job4.getId()), eq("batch-req-2"), eq(EmbeddingPrefixMode.OFF), eq(null));
        verify(executionPersistenceService).completeJob(eq(embeddingId5), aryEq(new float[]{1.3f, 1.4f, 1.5f}), eq(null), eq(job5.getId()), eq("batch-req-3"), eq(EmbeddingPrefixMode.OFF), eq(null));
        verify(executionPersistenceService).completeJob(eq(embeddingId6), aryEq(new float[]{1.6f, 1.7f, 1.8f}), eq(null), eq(job6.getId()), eq("batch-req-3"), eq(EmbeddingPrefixMode.OFF), eq(null));
    }
}
