package org.raddatz.familienarchiv.ocr;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.Loader;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.raddatz.familienarchiv.document.Document;
import org.raddatz.familienarchiv.document.annotation.DocumentAnnotation;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlock;
import org.raddatz.familienarchiv.document.annotation.AnnotationService;
import org.raddatz.familienarchiv.document.DocumentService;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlockQueryService;
import org.raddatz.familienarchiv.filestorage.FileService;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

@Service
@RequiredArgsConstructor
@Slf4j
public class TrainingDataExportService {

    private final TranscriptionBlockQueryService transcriptionBlockQueryService;
    private final AnnotationService annotationService;
    private final DocumentService documentService;
    private final FileService fileService;

    public List<TranscriptionBlock> queryEligibleBlocks() {
        return transcriptionBlockQueryService.findEligibleKurrentBlocks();
    }

    public StreamingResponseBody exportToZip() {
        return exportBlocksToZip(queryEligibleBlocks());
    }

    public List<TranscriptionBlock> queryBlocksForSender(UUID personId) {
        return transcriptionBlockQueryService.findManualKurrentBlocksByPerson(personId);
    }

    public StreamingResponseBody exportForSender(UUID personId) {
        return exportBlocksToZip(queryBlocksForSender(personId));
    }

    private StreamingResponseBody exportBlocksToZip(List<TranscriptionBlock> blocks) {
        if (blocks.isEmpty()) {
            return out -> {};
        }

        // Group blocks by documentId so we only download each PDF once
        Map<UUID, List<TranscriptionBlock>> byDoc = new LinkedHashMap<>();
        for (TranscriptionBlock b : blocks) {
            byDoc.computeIfAbsent(b.getDocumentId(), k -> new ArrayList<>()).add(b);
        }

        // Pre-fetch annotations keyed by id
        Map<UUID, DocumentAnnotation> annotations = new HashMap<>();
        for (TranscriptionBlock b : blocks) {
            annotationService.findById(b.getAnnotationId())
                    .ifPresent(a -> annotations.put(a.getId(), a));
        }

        // Pre-fetch documents keyed by id
        Map<UUID, Document> documents = new HashMap<>();
        for (UUID docId : byDoc.keySet()) {
            documentService.findById(docId).ifPresent(d -> documents.put(d.getId(), d));
        }

        return out -> {
            try (ZipOutputStream zip = new ZipOutputStream(out)) {
                for (Map.Entry<UUID, List<TranscriptionBlock>> entry : byDoc.entrySet()) {
                    UUID docId = entry.getKey();
                    Document doc = documents.get(docId);
                    if (doc == null || doc.getFilePath() == null) {
                        log.warn("Skipping document {} — no file path", docId);
                        continue;
                    }

                    byte[] pdfBytes;
                    try {
                        pdfBytes = fileService.downloadFileBytes(doc.getFilePath());
                    } catch (FileService.StorageFileNotFoundException | IOException e) {
                        log.warn("Skipping document {} — S3 download failed: {}", docId, e.getMessage());
                        continue;
                    }

                    try (PDDocument pdf = Loader.loadPDF(pdfBytes)) {
                        PDFRenderer renderer = new PDFRenderer(pdf);
                        for (TranscriptionBlock block : entry.getValue()) {
                            DocumentAnnotation ann = annotations.get(block.getAnnotationId());
                            if (ann == null) continue;

                            int pageIdx = ann.getPageNumber() - 1; // pageNumber is 1-based
                            if (pageIdx < 0 || pageIdx >= pdf.getNumberOfPages()) continue;

                            BufferedImage pageImage = renderPageImage(renderer, pageIdx);
                            BufferedImage cropped = cropBlockImage(pageImage, ann);

                            writeTrainingPair(zip, block.getId(), cropped, block.getText());
                        }
                    } catch (Exception e) {
                        log.warn("Skipping document {} — rendering failed: {}", docId, e.getMessage());
                    }
                }
            }
        };
    }

    BufferedImage renderPageImage(PDFRenderer renderer, int pageIdx) throws IOException {
        return renderer.renderImageWithDPI(pageIdx, 300);
    }

    BufferedImage cropBlockImage(BufferedImage page, DocumentAnnotation ann) {
        int imgW = page.getWidth();
        int imgH = page.getHeight();

        int x = (int) (ann.getX() * imgW);
        int y = (int) (ann.getY() * imgH);
        int w = (int) (ann.getWidth() * imgW);
        int h = (int) (ann.getHeight() * imgH);

        // Clamp to image bounds
        x = Math.max(0, Math.min(x, imgW - 1));
        y = Math.max(0, Math.min(y, imgH - 1));
        w = Math.max(1, Math.min(w, imgW - x));
        h = Math.max(1, Math.min(h, imgH - y));

        return page.getSubimage(x, y, w, h);
    }

    void writeTrainingPair(ZipOutputStream zip, UUID blockId, BufferedImage image, String text) throws IOException {
        String base = blockId.toString();
        int w = image.getWidth();
        int h = image.getHeight();
        // Baseline at 75 % height — typical text baseline position in a cropped line image
        int baselineY = (h * 3) / 4;

        // Write PNG
        zip.putNextEntry(new ZipEntry(base + ".png"));
        ImageIO.write(image, "PNG", zip);
        zip.closeEntry();

        // Write PAGE XML (Kraken 7+ dropped the legacy "path" format)
        String safeText = escapeXml(text != null ? text : "");
        String xml = String.format(
                "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n" +
                "<PcGts xmlns=\"http://schema.primaresearch.org/PAGE/gts/pagecontent/2019-07-15\">\n" +
                "  <Metadata><Creator>familienarchiv</Creator></Metadata>\n" +
                "  <Page imageFilename=\"%s.png\" imageWidth=\"%d\" imageHeight=\"%d\">\n" +
                "    <TextRegion id=\"r0\" type=\"paragraph\">\n" +
                "      <Coords points=\"0,0 %d,0 %d,%d 0,%d\"/>\n" +
                "      <TextLine id=\"l0\">\n" +
                "        <Coords points=\"0,0 %d,0 %d,%d 0,%d\"/>\n" +
                "        <Baseline points=\"0,%d %d,%d\"/>\n" +
                "        <TextEquiv><Unicode>%s</Unicode></TextEquiv>\n" +
                "      </TextLine>\n" +
                "    </TextRegion>\n" +
                "  </Page>\n" +
                "</PcGts>\n",
                base, w, h,
                w - 1, w - 1, h - 1, h - 1,
                w - 1, w - 1, h - 1, h - 1,
                baselineY, w - 1, baselineY,
                safeText);

        zip.putNextEntry(new ZipEntry(base + ".xml"));
        zip.write(xml.getBytes(StandardCharsets.UTF_8));
        zip.closeEntry();
    }

    private static String escapeXml(String text) {
        return text.replace("&", "&amp;")
                   .replace("<", "&lt;")
                   .replace(">", "&gt;");
    }
}
