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.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

@Service
@RequiredArgsConstructor
@Slf4j
public class SegmentationTrainingExportService {

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

    public List<TranscriptionBlock> querySegmentationBlocks() {
        return transcriptionBlockQueryService.findSegmentationBlocks();
    }

    public StreamingResponseBody exportToZip() {
        List<TranscriptionBlock> blocks = querySegmentationBlocks();
        if (blocks.isEmpty()) {
            return out -> {};
        }

        // Group by documentId so we download each PDF only 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;
                    }

                    // Group blocks by page number for this document
                    Map<Integer, List<TranscriptionBlock>> byPage = new LinkedHashMap<>();
                    for (TranscriptionBlock b : entry.getValue()) {
                        DocumentAnnotation ann = annotations.get(b.getAnnotationId());
                        if (ann != null) {
                            byPage.computeIfAbsent(ann.getPageNumber(), k -> new ArrayList<>()).add(b);
                        }
                    }

                    try (PDDocument pdf = Loader.loadPDF(pdfBytes)) {
                        PDFRenderer renderer = new PDFRenderer(pdf);
                        for (Map.Entry<Integer, List<TranscriptionBlock>> pageEntry : byPage.entrySet()) {
                            int pageNumber = pageEntry.getKey();
                            int pageIdx = pageNumber - 1;
                            if (pageIdx < 0 || pageIdx >= pdf.getNumberOfPages()) continue;

                            BufferedImage pageImage = renderer.renderImageWithDPI(pageIdx, 300);
                            String basename = "page-" + docId + "-" + pageNumber;

                            // Collect annotations for this page
                            List<DocumentAnnotation> pageAnnotations = new ArrayList<>();
                            for (TranscriptionBlock b : pageEntry.getValue()) {
                                DocumentAnnotation ann = annotations.get(b.getAnnotationId());
                                if (ann != null) pageAnnotations.add(ann);
                            }

                            writePngEntry(zip, basename, pageImage);
                            writePageXmlEntry(zip, basename, pageImage, pageAnnotations);
                        }
                    } catch (Exception e) {
                        log.warn("Skipping document {} — rendering failed: {}", docId, e.getMessage());
                    }
                }
            }
        };
    }

    private void writePngEntry(ZipOutputStream zip, String basename, BufferedImage image) throws IOException {
        zip.putNextEntry(new ZipEntry(basename + ".png"));
        ImageIO.write(image, "PNG", zip);
        zip.closeEntry();
    }

    private void writePageXmlEntry(ZipOutputStream zip, String basename,
                                   BufferedImage pageImage,
                                   List<DocumentAnnotation> annotations) throws IOException {
        int imgW = pageImage.getWidth();
        int imgH = pageImage.getHeight();

        StringBuilder regions = new StringBuilder();
        for (DocumentAnnotation ann : annotations) {
            String coords = buildPolygonCoords(ann, imgW, imgH);
            String regionId = ann.getId().toString();
            regions.append("      <TextRegion id=\"").append(regionId).append("\">\n");
            regions.append("        <Coords points=\"").append(coords).append("\"/>\n");
            regions.append("      </TextRegion>\n");
        }

        String xml = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
                + "<PcGts xmlns=\"http://schema.primaresearch.org/PAGE/gts/pagecontent/2019-07-15\">\n"
                + "  <Page imageFilename=\"" + basename + ".png\""
                + " imageWidth=\"" + imgW + "\""
                + " imageHeight=\"" + imgH + "\">\n"
                + regions
                + "  </Page>\n"
                + "</PcGts>\n";

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

    String buildPolygonCoords(DocumentAnnotation ann, int imgW, int imgH) {
        List<List<Double>> polygon = ann.getPolygon();
        if (polygon != null && !polygon.isEmpty()) {
            // Use explicit polygon — de-normalize to pixel coordinates
            StringBuilder sb = new StringBuilder();
            for (List<Double> pt : polygon) {
                if (sb.length() > 0) sb.append(' ');
                int px = (int) (pt.get(0) * imgW);
                int py = (int) (pt.get(1) * imgH);
                sb.append(px).append(',').append(py);
            }
            return sb.toString();
        }
        // Fall back to bounding box from x/y/width/height
        int x = (int) (ann.getX() * imgW);
        int y = (int) (ann.getY() * imgH);
        int w = (int) (ann.getWidth() * imgW);
        int h = (int) (ann.getHeight() * imgH);
        return x + "," + y + " " + (x + w) + "," + y + " " + (x + w) + "," + (y + h) + " " + x + "," + (y + h);
    }
}
