package org.raddatz.familienarchiv.ocr;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.raddatz.familienarchiv.document.annotation.CreateAnnotationDTO;
import org.raddatz.familienarchiv.document.annotation.DocumentAnnotation;
import org.raddatz.familienarchiv.document.Document;
import org.raddatz.familienarchiv.document.DocumentStatus;
import org.raddatz.familienarchiv.ocr.OcrJobDocumentRepository;
import org.raddatz.familienarchiv.ocr.OcrJobRepository;

import org.raddatz.familienarchiv.document.DocumentService;
import org.raddatz.familienarchiv.document.transcription.TranscriptionService;
import org.raddatz.familienarchiv.document.annotation.AnnotationService;
import org.raddatz.familienarchiv.filestorage.FileService;
import org.raddatz.familienarchiv.ocr.SenderModelService;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import java.util.function.Consumer;

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

@ExtendWith(MockitoExtension.class)
class OcrAsyncRunnerTest {

    @Mock OcrClient ocrClient;
    @Mock DocumentService documentService;
    @Mock TranscriptionService transcriptionService;
    @Mock AnnotationService annotationService;
    @Mock FileService fileService;
    @Mock OcrJobRepository ocrJobRepository;
    @Mock OcrJobDocumentRepository ocrJobDocumentRepository;
    @Mock OcrProgressService ocrProgressService;
    @Mock SenderModelService senderModelService;

    @InjectMocks OcrAsyncRunner ocrAsyncRunner;

    @Test
    void processDocument_clearsExistingBlocks() {
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();

        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");
        doAnswer(inv -> {
            Consumer<OcrStreamEvent> handler = inv.getArgument(4);
            handler.accept(new OcrStreamEvent.Start(0));
            handler.accept(new OcrStreamEvent.Done(0, 0));
            return null;
        }).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());

        ocrAsyncRunner.processDocument(docId, doc, userId);

        verify(transcriptionService).deleteAllBlocksByDocument(docId);
    }

    @Test
    void processDocument_createsAnnotationAndBlock_forEachResult() {
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        UUID annId = UUID.randomUUID();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();


        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");
        doAnswer(inv -> {
            Consumer<OcrStreamEvent> handler = inv.getArgument(4);
            handler.accept(new OcrStreamEvent.Start(1));
            handler.accept(new OcrStreamEvent.Page(0, List.of(
                    new OcrBlockResult(0, 0.1, 0.1, 0.8, 0.04, null, "Line 1", null),
                    new OcrBlockResult(0, 0.1, 0.2, 0.8, 0.04, null, "Line 2", null))));
            handler.accept(new OcrStreamEvent.Done(2, 0));
            return null;
        }).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());
        DocumentAnnotation ann = DocumentAnnotation.builder().id(annId).build();
        when(annotationService.createOcrAnnotation(any(), any(), any(), any(), any())).thenReturn(ann);

        ocrAsyncRunner.processDocument(docId, doc, userId);

        verify(annotationService, times(2)).createOcrAnnotation(
                eq(docId), any(CreateAnnotationDTO.class), eq(userId), eq("hash"), any());
        verify(transcriptionService, times(2)).createOcrBlock(
                eq(docId), eq(annId), any(), anyInt(), eq(userId));
    }

    @Test
    void processDocument_delegatesBlockCreationToTranscriptionService() {
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        UUID annId = UUID.randomUUID();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();


        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");
        doAnswer(inv -> {
            Consumer<OcrStreamEvent> handler = inv.getArgument(4);
            handler.accept(new OcrStreamEvent.Start(1));
            handler.accept(new OcrStreamEvent.Page(0, List.of(
                    new OcrBlockResult(0, 0.1, 0.1, 0.8, 0.04, null, "Test", null))));
            handler.accept(new OcrStreamEvent.Done(1, 0));
            return null;
        }).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());
        DocumentAnnotation ann = DocumentAnnotation.builder().id(annId).build();
        when(annotationService.createOcrAnnotation(any(), any(), any(), any(), any())).thenReturn(ann);

        ocrAsyncRunner.processDocument(docId, doc, userId);

        verify(transcriptionService).createOcrBlock(docId, annId, "Test", 0, userId);
    }

    @Test
    void runSingleDocument_setsJobDone_onSuccess() {
        UUID jobId = UUID.randomUUID();
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        OcrJob job = OcrJob.builder().id(jobId).totalDocuments(1).status(OcrJobStatus.PENDING).build();
        OcrJobDocument jobDoc = OcrJobDocument.builder().id(UUID.randomUUID())
                .jobId(jobId).documentId(docId).status(OcrDocumentStatus.PENDING).build();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();

        when(ocrJobRepository.findById(jobId)).thenReturn(Optional.of(job));
        when(ocrJobRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(ocrJobDocumentRepository.findByJobIdAndDocumentId(jobId, docId))
                .thenReturn(Optional.of(jobDoc));
        when(documentService.getDocumentById(docId)).thenReturn(doc);

        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");
        doAnswer(inv -> {
            Consumer<OcrStreamEvent> handler = inv.getArgument(4);
            handler.accept(new OcrStreamEvent.Start(1));
            handler.accept(new OcrStreamEvent.Page(0, List.of()));
            handler.accept(new OcrStreamEvent.Done(0, 0));
            return null;
        }).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());

        ocrAsyncRunner.runSingleDocument(jobId, docId, userId);

        assertThat(job.getStatus()).isEqualTo(OcrJobStatus.DONE);
    }

    @Test
    void runSingleDocument_setsJobFailed_onError() {
        UUID jobId = UUID.randomUUID();
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        OcrJob job = OcrJob.builder().id(jobId).totalDocuments(1).status(OcrJobStatus.PENDING).build();
        OcrJobDocument jobDoc = OcrJobDocument.builder().id(UUID.randomUUID())
                .jobId(jobId).documentId(docId).status(OcrDocumentStatus.PENDING).build();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();

        when(ocrJobRepository.findById(jobId)).thenReturn(Optional.of(job));
        when(ocrJobRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(ocrJobDocumentRepository.findByJobIdAndDocumentId(jobId, docId))
                .thenReturn(Optional.of(jobDoc));
        when(documentService.getDocumentById(docId)).thenReturn(doc);

        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");
        doThrow(new RuntimeException("OCR failed")).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());

        ocrAsyncRunner.runSingleDocument(jobId, docId, userId);

        assertThat(job.getStatus()).isEqualTo(OcrJobStatus.FAILED);
        assertThat(job.getErrorCount()).isEqualTo(1);
    }

    @Test
    void runSingleDocument_updatesProgressPerPage() {
        UUID jobId = UUID.randomUUID();
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        OcrJob job = OcrJob.builder().id(jobId).totalDocuments(1).status(OcrJobStatus.PENDING).build();
        OcrJobDocument jobDoc = OcrJobDocument.builder().id(UUID.randomUUID())
                .jobId(jobId).documentId(docId).status(OcrDocumentStatus.PENDING).build();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();
        DocumentAnnotation ann = DocumentAnnotation.builder().id(UUID.randomUUID()).build();

        when(ocrJobRepository.findById(jobId)).thenReturn(Optional.of(job));
        when(ocrJobRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(ocrJobDocumentRepository.findByJobIdAndDocumentId(jobId, docId))
                .thenReturn(Optional.of(jobDoc));
        when(ocrJobDocumentRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(documentService.getDocumentById(docId)).thenReturn(doc);

        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");
        when(annotationService.createOcrAnnotation(any(), any(), any(), any(), any())).thenReturn(ann);

        List<String> progressMessages = new ArrayList<>();
        doAnswer(inv -> {
            Consumer<OcrStreamEvent> handler = inv.getArgument(4);
            handler.accept(new OcrStreamEvent.Start(3));
            handler.accept(new OcrStreamEvent.Page(0, List.of(
                    new OcrBlockResult(0, 0.1, 0.1, 0.8, 0.04, null, "L1", null),
                    new OcrBlockResult(0, 0.1, 0.2, 0.8, 0.04, null, "L2", null))));
            progressMessages.add(job.getProgressMessage());
            handler.accept(new OcrStreamEvent.Page(1, List.of(
                    new OcrBlockResult(1, 0.1, 0.1, 0.8, 0.04, null, "L3", null))));
            progressMessages.add(job.getProgressMessage());
            handler.accept(new OcrStreamEvent.Done(3, 0));
            return null;
        }).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());

        ocrAsyncRunner.runSingleDocument(jobId, docId, userId);

        assertThat(progressMessages.get(0)).isEqualTo("ANALYZING_PAGE:1:3:2");
        assertThat(progressMessages.get(1)).isEqualTo("ANALYZING_PAGE:2:3:3");
        assertThat(job.getProgressMessage()).isEqualTo("DONE:3:0");
    }

    @Test
    void runSingleDocument_includesSkippedPagesInDoneMessage() {
        UUID jobId = UUID.randomUUID();
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        OcrJob job = OcrJob.builder().id(jobId).totalDocuments(1).status(OcrJobStatus.PENDING).build();
        OcrJobDocument jobDoc = OcrJobDocument.builder().id(UUID.randomUUID())
                .jobId(jobId).documentId(docId).status(OcrDocumentStatus.PENDING).build();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();

        when(ocrJobRepository.findById(jobId)).thenReturn(Optional.of(job));
        when(ocrJobRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(ocrJobDocumentRepository.findByJobIdAndDocumentId(jobId, docId))
                .thenReturn(Optional.of(jobDoc));
        when(ocrJobDocumentRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(documentService.getDocumentById(docId)).thenReturn(doc);

        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");

        doAnswer(inv -> {
            Consumer<OcrStreamEvent> handler = inv.getArgument(4);
            handler.accept(new OcrStreamEvent.Start(3));
            handler.accept(new OcrStreamEvent.Page(0, List.of()));
            handler.accept(new OcrStreamEvent.Error(1, "failed"));
            handler.accept(new OcrStreamEvent.Page(2, List.of()));
            handler.accept(new OcrStreamEvent.Done(0, 1));
            return null;
        }).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());

        ocrAsyncRunner.runSingleDocument(jobId, docId, userId);

        assertThat(job.getStatus()).isEqualTo(OcrJobStatus.DONE);
        assertThat(job.getProgressMessage()).isEqualTo("DONE:0:1");
    }

    @Test
    void runSingleDocument_updatesProgressOnPreprocessingEvent() {
        UUID jobId = UUID.randomUUID();
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        OcrJob job = OcrJob.builder().id(jobId).totalDocuments(1).status(OcrJobStatus.PENDING).build();
        OcrJobDocument jobDoc = OcrJobDocument.builder().id(UUID.randomUUID())
                .jobId(jobId).documentId(docId).status(OcrDocumentStatus.PENDING).build();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();

        when(ocrJobRepository.findById(jobId)).thenReturn(Optional.of(job));
        when(ocrJobRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(ocrJobDocumentRepository.findByJobIdAndDocumentId(jobId, docId))
                .thenReturn(Optional.of(jobDoc));
        when(ocrJobDocumentRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(documentService.getDocumentById(docId)).thenReturn(doc);
        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");

        List<String> progressMessages = new ArrayList<>();
        doAnswer(inv -> {
            Consumer<OcrStreamEvent> handler = inv.getArgument(4);
            handler.accept(new OcrStreamEvent.Start(5));
            handler.accept(new OcrStreamEvent.Preprocessing(1));
            progressMessages.add(job.getProgressMessage());
            handler.accept(new OcrStreamEvent.Page(1, List.of()));
            handler.accept(new OcrStreamEvent.Done(0, 0));
            return null;
        }).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());

        ocrAsyncRunner.runSingleDocument(jobId, docId, userId);

        // Must have totalPages=5 from the preceding Start event, not 0
        assertThat(progressMessages.get(0)).isEqualTo("PREPROCESSING_PAGE:1:5");
    }

    @Test
    void runSingleDocument_logsStreamErrorAtWarnWithoutSettingJobFailed() {
        UUID jobId = UUID.randomUUID();
        UUID docId = UUID.randomUUID();
        UUID userId = UUID.randomUUID();
        OcrJob job = OcrJob.builder().id(jobId).totalDocuments(1).status(OcrJobStatus.PENDING).build();
        OcrJobDocument jobDoc = OcrJobDocument.builder().id(UUID.randomUUID())
                .jobId(jobId).documentId(docId).status(OcrDocumentStatus.PENDING).build();
        Document doc = Document.builder().id(docId).filePath("test.pdf")
                .fileHash("hash").scriptType(ScriptType.TYPEWRITER).build();

        when(ocrJobRepository.findById(jobId)).thenReturn(Optional.of(job));
        when(ocrJobRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(ocrJobDocumentRepository.findByJobIdAndDocumentId(jobId, docId))
                .thenReturn(Optional.of(jobDoc));
        when(ocrJobDocumentRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(documentService.getDocumentById(docId)).thenReturn(doc);

        when(fileService.generatePresignedUrl(any())).thenReturn("http://presigned");

        doAnswer(inv -> {
            Consumer<OcrStreamEvent> handler = inv.getArgument(4);
            handler.accept(new OcrStreamEvent.Start(2));
            handler.accept(new OcrStreamEvent.Error(0, "some python traceback details"));
            handler.accept(new OcrStreamEvent.Page(1, List.of()));
            handler.accept(new OcrStreamEvent.Done(0, 1));
            return null;
        }).when(ocrClient).streamBlocks(any(), any(), any(), any(), any());

        ocrAsyncRunner.runSingleDocument(jobId, docId, userId);

        // Job should still be DONE, not FAILED (per-page errors don't fail the whole job)
        assertThat(job.getStatus()).isEqualTo(OcrJobStatus.DONE);
        // Raw error message should not leak to progress
        assertThat(job.getProgressMessage()).doesNotContain("python traceback");
    }
}
