package org.raddatz.familienarchiv.ocr;

import org.apache.pdfbox.Loader;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.pdmodel.common.PDRectangle;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.raddatz.familienarchiv.PostgresContainerConfig;
import org.raddatz.familienarchiv.config.FlywayConfig;
import org.raddatz.familienarchiv.document.annotation.AnnotationRepository;
import org.raddatz.familienarchiv.document.annotation.DocumentAnnotation;
import org.raddatz.familienarchiv.document.BlockSource;
import org.raddatz.familienarchiv.document.Document;
import org.raddatz.familienarchiv.document.DocumentRepository;
import org.raddatz.familienarchiv.document.DocumentStatus;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlock;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlockRepository;

import org.raddatz.familienarchiv.document.DocumentService;
import org.raddatz.familienarchiv.document.annotation.AnnotationService;
import org.raddatz.familienarchiv.filestorage.FileService;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlockQueryService;
import org.raddatz.familienarchiv.ocr.TrainingLabel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.jdbc.test.autoconfigure.AutoConfigureTestDatabase;
import org.springframework.boot.data.jpa.test.autoconfigure.DataJpaTest;
import org.springframework.context.annotation.Import;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;

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

@DataJpaTest
@AutoConfigureTestDatabase(replace = AutoConfigureTestDatabase.Replace.NONE)
@Import({PostgresContainerConfig.class, FlywayConfig.class})
class TrainingDataExportServiceTest {

    @Autowired TranscriptionBlockRepository blockRepository;
    @Autowired DocumentRepository documentRepository;
    @Autowired AnnotationRepository annotationRepository;

    static byte[] minimalPdfBytes;

    @BeforeAll
    static void createMinimalPdf() throws Exception {
        try (PDDocument doc = new PDDocument()) {
            doc.addPage(new PDPage(PDRectangle.A4));
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            doc.save(out);
            minimalPdfBytes = out.toByteArray();
        }
    }

    // ─── Query: enrollment filter ─────────────────────────────────────────────

    @Test
    void export_includesManualBlockFromEnrolledDocument() throws Exception {
        UUID docId = enrolledDoc("enrolled.pdf");
        UUID annotId = annotation(docId);
        blockRepository.save(manualBlock(docId, annotId, "Liebe Mutter"));

        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        StreamingResponseBody body = service.exportToZip();
        byte[] zipBytes = stream(body);
        assertThat(zipEntryNames(zipBytes)).isNotEmpty();
    }

    @Test
    void export_excludesUnreviewedManualBlockFromEnrolledDocument() throws Exception {
        // MANUAL blocks whose text hasn't been verified yet should not count as training data
        UUID docId = enrolledDoc("unreviewed-manual.pdf");
        UUID annotId = annotation(docId);
        TranscriptionBlock block = TranscriptionBlock.builder()
                .annotationId(annotId).documentId(docId)
                .text("Liebe Tante").sortOrder(0)
                .source(BlockSource.MANUAL).reviewed(false).build();
        blockRepository.save(block);

        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        StreamingResponseBody body = service.exportToZip();
        assertThat(zipEntryNames(stream(body))).isEmpty();
    }

    @Test
    void export_excludesManualBlockFromNonEnrolledDocument() throws Exception {
        UUID docId = nonEnrolledDoc("notenrolled.pdf");
        UUID annotId = annotation(docId);
        blockRepository.save(manualBlock(docId, annotId, "Liebe Tante"));

        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        StreamingResponseBody body = service.exportToZip();
        byte[] zipBytes = stream(body);
        assertThat(zipEntryNames(zipBytes)).isEmpty();
    }

    @Test
    void export_includesReviewedOcrBlockFromEnrolledDocument() throws Exception {
        UUID docId = enrolledDoc("ocr-reviewed.pdf");
        UUID annotId = annotation(docId);
        TranscriptionBlock block = TranscriptionBlock.builder()
                .annotationId(annotId).documentId(docId)
                .text("OCR text").sortOrder(0)
                .source(BlockSource.OCR).reviewed(true).build();
        blockRepository.save(block);

        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        StreamingResponseBody body = service.exportToZip();
        assertThat(zipEntryNames(stream(body))).isNotEmpty();
    }

    @Test
    void export_excludesUnreviewedOcrBlockFromEnrolledDocument() throws Exception {
        UUID docId = enrolledDoc("ocr-unreviewed.pdf");
        UUID annotId = annotation(docId);
        TranscriptionBlock block = TranscriptionBlock.builder()
                .annotationId(annotId).documentId(docId)
                .text("Raw OCR").sortOrder(0)
                .source(BlockSource.OCR).reviewed(false).build();
        blockRepository.save(block);

        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        StreamingResponseBody body = service.exportToZip();
        assertThat(zipEntryNames(stream(body))).isEmpty();
    }

    // ─── ZIP structure ────────────────────────────────────────────────────────

    @Test
    void export_producesExactly2EntriesPerBlock_pngAndTxt() throws Exception {
        UUID docId = enrolledDoc("zip-struct.pdf");
        UUID annotId = annotation(docId);
        blockRepository.save(manualBlock(docId, annotId, "Erste Zeile"));
        blockRepository.save(manualBlock(docId, annotId, "Zweite Zeile"));

        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        byte[] zipBytes = stream(service.exportToZip());
        var names = zipEntryNames(zipBytes);
        assertThat(names).hasSize(4); // 2 blocks × 2 entries each
        assertThat(names.stream().filter(n -> n.endsWith(".png")).count()).isEqualTo(2);
        assertThat(names.stream().filter(n -> n.endsWith(".xml")).count()).isEqualTo(2);
    }

    @Test
    void export_pageXmlContainsBlockText() throws Exception {
        UUID docId = enrolledDoc("txt-content.pdf");
        UUID annotId = annotation(docId);
        String expectedText = "Sehr geehrte Frau";
        blockRepository.save(manualBlock(docId, annotId, expectedText));

        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        byte[] zipBytes = stream(service.exportToZip());
        String xmlContent = readZipEntry(zipBytes, ".xml");
        assertThat(xmlContent).contains("<Unicode>" + expectedText + "</Unicode>");
    }

    @Test
    void export_pageXmlEscapesSpecialCharacters() throws Exception {
        UUID docId = enrolledDoc("special-chars.pdf");
        UUID annotId = annotation(docId);
        blockRepository.save(manualBlock(docId, annotId, "A & B < C > D"));

        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        byte[] zipBytes = stream(service.exportToZip());
        String xmlContent = readZipEntry(zipBytes, ".xml");
        assertThat(xmlContent).contains("<Unicode>A &amp; B &lt; C &gt; D</Unicode>");
    }

    // ─── S3 failure resilience ────────────────────────────────────────────────

    @Test
    void export_skipsDocumentWhenS3DownloadFails_andStillIncludesOtherDocuments() throws Exception {
        UUID failDocId = enrolledDoc("fail.pdf");
        UUID okDocId = enrolledDoc("ok.pdf");
        UUID failAnnotId = annotation(failDocId);
        UUID okAnnotId = annotation(okDocId);
        blockRepository.save(manualBlock(failDocId, failAnnotId, "Will fail"));
        blockRepository.save(manualBlock(okDocId, okAnnotId, "Will succeed"));

        FileService fileService = mock(FileService.class);
        when(fileService.downloadFileBytes("fail.pdf")).thenThrow(new FileService.StorageFileNotFoundException("missing"));
        when(fileService.downloadFileBytes("ok.pdf")).thenReturn(minimalPdfBytes);

        TrainingDataExportService service = makeService(fileService);

        byte[] zipBytes = stream(service.exportToZip());
        var names = zipEntryNames(zipBytes);
        // ok.pdf block produces 2 entries; fail.pdf block is skipped
        assertThat(names).hasSize(2);
    }

    // ─── Empty export ─────────────────────────────────────────────────────────

    @Test
    void queryEligibleBlocks_returnsEmpty_whenNoEnrolledDocuments() {
        FileService fileService = mockFileService();
        TrainingDataExportService service = makeService(fileService);

        assertThat(service.queryEligibleBlocks()).isEmpty();
    }

    // ─── helpers ─────────────────────────────────────────────────────────────

    /**
     * Builds the export service with mocked owning services that transparently
     * delegate every read to the real JPA repositories provided by {@code @DataJpaTest}.
     * Keeps real-database fidelity without pulling the full service trees into scope.
     */
    private TrainingDataExportService makeService(FileService fileService) {
        TranscriptionBlockQueryService blockQueryService = mock(TranscriptionBlockQueryService.class);
        AnnotationService annotationService = mock(AnnotationService.class);
        DocumentService documentService = mock(DocumentService.class);
        when(blockQueryService.findEligibleKurrentBlocks())
                .thenAnswer(inv -> blockRepository.findEligibleKurrentBlocks());
        when(blockQueryService.findManualKurrentBlocksByPerson(any(UUID.class)))
                .thenAnswer(inv -> blockRepository.findManualKurrentBlocksByPerson(inv.getArgument(0)));
        when(annotationService.findById(any(UUID.class)))
                .thenAnswer(inv -> annotationRepository.findById(inv.getArgument(0)));
        when(documentService.findById(any(UUID.class)))
                .thenAnswer(inv -> documentRepository.findById(inv.getArgument(0)));
        return new TrainingDataExportService(blockQueryService, annotationService, documentService, fileService);
    }

    private UUID enrolledDoc(String filename) {
        Document doc = documentRepository.save(Document.builder()
                .title(filename).originalFilename(filename).filePath(filename)
                .status(DocumentStatus.UPLOADED)
                .trainingLabels(new HashSet<>(Set.of(TrainingLabel.KURRENT_RECOGNITION)))
                .build());
        return doc.getId();
    }

    private UUID nonEnrolledDoc(String filename) {
        Document doc = documentRepository.save(Document.builder()
                .title(filename).originalFilename(filename).filePath(filename)
                .status(DocumentStatus.UPLOADED)
                .build());
        return doc.getId();
    }

    private UUID annotation(UUID docId) {
        return annotationRepository.save(DocumentAnnotation.builder()
                .documentId(docId).pageNumber(1)
                .x(0.1).y(0.1).width(0.8).height(0.1).color("#00C7B1")
                .build()).getId();
    }

    private TranscriptionBlock manualBlock(UUID docId, UUID annotId, String text) {
        return TranscriptionBlock.builder()
                .annotationId(annotId).documentId(docId)
                .text(text).sortOrder(0)
                .source(BlockSource.MANUAL).reviewed(true).build();
    }

    private FileService mockFileService() {
        FileService fs = mock(FileService.class);
        try {
            when(fs.downloadFileBytes(anyString())).thenReturn(minimalPdfBytes);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return fs;
    }

    private static byte[] stream(StreamingResponseBody body) throws Exception {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        body.writeTo(out);
        return out.toByteArray();
    }

    private static java.util.List<String> zipEntryNames(byte[] zipBytes) throws Exception {
        var names = new java.util.ArrayList<String>();
        try (ZipInputStream zis = new ZipInputStream(new ByteArrayInputStream(zipBytes))) {
            ZipEntry entry;
            while ((entry = zis.getNextEntry()) != null) {
                names.add(entry.getName());
                zis.closeEntry();
            }
        }
        return names;
    }

    private static String readZipEntry(byte[] zipBytes, String suffix) throws Exception {
        try (ZipInputStream zis = new ZipInputStream(new ByteArrayInputStream(zipBytes))) {
            ZipEntry entry;
            while ((entry = zis.getNextEntry()) != null) {
                if (entry.getName().endsWith(suffix)) {
                    return new String(zis.readAllBytes(), java.nio.charset.StandardCharsets.UTF_8);
                }
                zis.closeEntry();
            }
        }
        return null;
    }
}
