package com.knecon.fforesight.service.viewerdoc.service;

import java.awt.geom.AffineTransform;
import java.awt.geom.Rectangle2D;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

import org.apache.pdfbox.Loader;
import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.pdfwriter.compress.CompressParameters;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDDocumentCatalog;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.pdmodel.PDPageContentStream;
import org.apache.pdfbox.pdmodel.PDResources;
import org.apache.pdfbox.pdmodel.font.PDFont;
import org.apache.pdfbox.pdmodel.graphics.optionalcontent.PDOptionalContentGroup;
import org.apache.pdfbox.pdmodel.graphics.optionalcontent.PDOptionalContentProperties;
import org.apache.pdfbox.pdmodel.graphics.state.PDExtendedGraphicsState;
import org.apache.pdfbox.pdmodel.graphics.state.RenderingMode;
import org.apache.pdfbox.util.Matrix;

import com.knecon.fforesight.service.viewerdoc.ContentStreams;
import com.knecon.fforesight.service.viewerdoc.model.ColoredLine;
import com.knecon.fforesight.service.viewerdoc.model.ColoredRectangle;
import com.knecon.fforesight.service.viewerdoc.model.FilledRectangle;
import com.knecon.fforesight.service.viewerdoc.model.PlacedText;
import com.knecon.fforesight.service.viewerdoc.model.Visualizations;
import com.knecon.fforesight.service.viewerdoc.model.VisualizationsOnPage;
import com.knecon.fforesight.service.viewerdoc.pdf.ClassifiedContentStream;

import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.annotation.Observed;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

@Slf4j
@RequiredArgsConstructor
public class ViewerDocumentService {

    private final ObservationRegistry registry;


    @Observed(name = "ViewerDocumentService", contextualName = "add-visualizations")
    @SneakyThrows
    public void addVisualizationsOnPage(File originFile, File destinationFile, List<Visualizations> visualizations) {

        // originFile and destinationFile might be the same, so we use a temp file.
        // Otherwise, saving the document might corrupt the file
        Path tmpFile = Files.createTempFile("tmpViewerDocument", ".pdf");
        Files.copy(originFile.toPath(), tmpFile, StandardCopyOption.REPLACE_EXISTING);

        PDDocument pdDocument = openPDDocument(tmpFile.toFile());

        enrichObservation(pdDocument,
                          visualizations.stream()
                                  .map(Visualizations::getLayer)
                                  .toList());

        Set<ContentStreams.Identifier> allLayers = visualizations.stream()
                .map(Visualizations::getLayer)
                .collect(Collectors.toUnmodifiableSet());

        Map<ContentStreams.Identifier, PDOptionalContentGroup> optionalContentGroupMap = addLayersToDocument(visualizations, pdDocument);

        for (int pageNumber = 0; pageNumber < pdDocument.getNumberOfPages(); pageNumber++) {
            PDPage pdPage = pdDocument.getPage(pageNumber);
            createPageResourcesIfNotPresent(pdPage); // needed for optionalContentGroups

            List<ClassifiedContentStream> classifiers = ContentStreamClassifier.getClassifiedContentStreams(pdPage);

            pdPage.setContents(ContentStreamUtility.removeLayerFromContentStreams(allLayers, classifiers));

            AffineTransform textDeRotationMatrix = getTextDeRotationTransform(pdPage);

            if (!ContentStreamClassifier.areAllContentStreamsEscaped(classifiers)) {
                // We need to save the graphics state before, such that our appended content cannot be affected by previous content streams with side effects,
                // e.g. not escaped matrix transformations.
                wrapContentStreams(pdDocument, pdPage);
            }

            for (Visualizations visualization : visualizations) {
                if (!visualization.getVisualizationsOnPages().containsKey(pageNumber)) {
                    continue;
                }
                // We need to append to the content stream, otherwise the content could be overlapped by following content.
                try (var contentStream = new PDPageContentStream(pdDocument, pdPage, PDPageContentStream.AppendMode.APPEND, true)) {

                    contentStream.beginMarkedContent(visualization.getLayer().cosName());

                    if (optionalContentGroupMap.containsKey(visualization.getLayer())) {
                        contentStream.beginMarkedContent(COSName.OC, optionalContentGroupMap.get(visualization.getLayer()));
                    }

                    contentStream.saveGraphicsState();

                    drawVisualizationsToContentStream(pdDocument, visualization.getVisualizationsOnPages().get(pageNumber), contentStream, textDeRotationMatrix);

                    contentStream.restoreGraphicsState();

                    if (optionalContentGroupMap.containsKey(visualization.getLayer())) {
                        contentStream.endMarkedContent();
                    }

                    contentStream.endMarkedContent();
                }

            }
            if (pageNumber % 500 == 0 && pageNumber != 0) { // re-open document every once in a while to save on RAM
                log.info("Incremental save after {}/{} pages", pageNumber, pdDocument.getNumberOfPages());
                observedIncrementalSave(pdDocument, destinationFile);
                pdDocument.close();
                Files.copy(destinationFile.toPath(), tmpFile, StandardCopyOption.REPLACE_EXISTING);
                pdDocument = openPDDocument(tmpFile.toFile());
            }
        }
        observedIncrementalSave(pdDocument, destinationFile);

        pdDocument.close();
        assert tmpFile.toFile().delete();
    }


    private static Map<ContentStreams.Identifier, PDOptionalContentGroup> addLayersToDocument(List<Visualizations> visualizations, PDDocument pdDocument) {

        Map<ContentStreams.Identifier, PDOptionalContentGroup> optionalContentGroupMap = new HashMap<>();
        for (Visualizations visualization : visualizations) {
            addLayerToDocument(visualization.getLayer(), pdDocument, visualization.isLayerVisibilityDefaultValue())//
                    .ifPresent(ocg -> optionalContentGroupMap.put(visualization.getLayer(), ocg));
        }
        return optionalContentGroupMap;
    }


    private static void drawVisualizationsToContentStream(PDDocument pdDocument,
                                                          VisualizationsOnPage visualizationsOnPage,
                                                          PDPageContentStream contentStream,
                                                          AffineTransform textDeRotationMatrix) throws IOException {

        if (visualizationsOnPage.isMakePathsInvisible()) {
            contentStream.addRect(0, 0, 1, 1);
            contentStream.clip();
        }

        for (ColoredLine coloredLine : visualizationsOnPage.getColoredLines()) {
            contentStream.setLineWidth(coloredLine.lineWidth());
            contentStream.setStrokingColor(coloredLine.color());
            contentStream.moveTo((float) coloredLine.line().getX1(), (float) coloredLine.line().getY1());
            contentStream.lineTo((float) coloredLine.line().getX2(), (float) coloredLine.line().getY2());
            contentStream.stroke();
        }

        for (ColoredRectangle coloredRectangle : visualizationsOnPage.getColoredRectangles()) {
            contentStream.setLineWidth(coloredRectangle.lineWidth());
            contentStream.setStrokingColor(coloredRectangle.color());
            Rectangle2D r = coloredRectangle.rectangle2D();
            contentStream.addRect((float) r.getX(), (float) r.getY(), (float) r.getWidth(), (float) r.getHeight());
            contentStream.stroke();
        }

        for (FilledRectangle filledRectangle : visualizationsOnPage.getFilledRectangles()) {
            contentStream.setNonStrokingColor(filledRectangle.color());
            PDExtendedGraphicsState graphicsState = new PDExtendedGraphicsState();
            graphicsState.setNonStrokingAlphaConstant(filledRectangle.alpha());
            contentStream.setGraphicsStateParameters(graphicsState);
            Rectangle2D r = filledRectangle.rectangle2D();
            contentStream.addRect((float) r.getX(), (float) r.getY(), (float) r.getWidth(), (float) r.getHeight());
            contentStream.fill();
        }

        for (PlacedText placedText : visualizationsOnPage.getPlacedTexts()) {
            PDFont font = placedText.font().embed(pdDocument);
            contentStream.setFont(font, placedText.fontSize());
            contentStream.beginText();
            contentStream.setNonStrokingColor(placedText.color());
            if (placedText.renderingMode()
                    .isPresent()) {
                contentStream.setRenderingMode(placedText.renderingMode().get());
            } else {
                contentStream.setRenderingMode(RenderingMode.FILL);
            }
            Matrix textMatrix = getTextMatrix(placedText, textDeRotationMatrix);
            contentStream.setTextMatrix(textMatrix);
            contentStream.showText(placedText.text());
            contentStream.endText();
        }
    }


    private void enrichObservation(PDDocument pdDocument, List<ContentStreams.Identifier> layers) {

        if (registry == null || registry.getCurrentObservation() == null || registry.isNoop()) {
            return;
        }
        registry.getCurrentObservation().highCardinalityKeyValue("numberOfPages", String.valueOf(pdDocument.getNumberOfPages()));
        for (int i = 0; i < layers.size(); i++) {
            ContentStreams.Identifier layer = layers.get(i);

            registry.getCurrentObservation().highCardinalityKeyValue("layer_" + i, String.valueOf(layer.name()));
        }
    }


    private static void wrapContentStreams(PDDocument pdDocument, PDPage pdPage) throws IOException {

        try (var contentStream = new PDPageContentStream(pdDocument, pdPage, PDPageContentStream.AppendMode.PREPEND, false)) {
            contentStream.saveGraphicsState();
        }
        try (var contentStream = new PDPageContentStream(pdDocument, pdPage, PDPageContentStream.AppendMode.APPEND, false)) {
            contentStream.restoreGraphicsState();
        }
    }


    private static Matrix getTextMatrix(PlacedText placedText, AffineTransform textDeRotationMatrix) {

        Matrix textMatrix;
        if (placedText.textMatrix().isEmpty()) {
            textMatrix = new Matrix((float) textDeRotationMatrix.getScaleX(),
                                    (float) textDeRotationMatrix.getShearX(),
                                    (float) textDeRotationMatrix.getShearY(),
                                    (float) textDeRotationMatrix.getScaleY(),
                                    (float) placedText.lineStart().getX(),
                                    (float) placedText.lineStart().getY());
        } else {
            textMatrix = placedText.textMatrix().get();
        }
        return textMatrix;
    }


    private static Optional<PDOptionalContentGroup> addLayerToDocument(ContentStreams.Identifier layer, PDDocument pdDocument, boolean layerVisibilityDefaultValue) {

        if (layer.optionalContent()) {
            return Optional.of(addLayerToDocument(pdDocument, layer.name(), layerVisibilityDefaultValue));

        }
        return Optional.empty();

    }


    private static PDOptionalContentGroup addLayerToDocument(PDDocument pdDocument, String layerName, boolean layerVisibilityDefaultValue) {

        PDDocumentCatalog catalog = pdDocument.getDocumentCatalog();
        PDOptionalContentProperties ocprops = catalog.getOCProperties();
        if (ocprops == null) {
            ocprops = new PDOptionalContentProperties();
            catalog.setOCProperties(ocprops);
        }
        PDOptionalContentGroup layer = null;
        if (ocprops.hasGroup(layerName)) {
            layer = ocprops.getGroup(layerName);
        } else {
            layer = new PDOptionalContentGroup(layerName);
            ocprops.addGroup(layer);
        }
        ocprops.setGroupEnabled(layer, layerVisibilityDefaultValue);
        return layer;
    }


    private static PDDocument openPDDocument(File tmpFile) throws IOException {

        PDDocument pdDocument;
        pdDocument = Loader.loadPDF(tmpFile);
        pdDocument.setAllSecurityToBeRemoved(true);
        return pdDocument;
    }


    @SneakyThrows
    private void observedIncrementalSave(PDDocument pdDocument, File outputFile) {

        Observation.createNotStarted("ViewerDocumentService", registry).contextualName("incremental-save").observe(() -> {
            try (var out = new FileOutputStream(outputFile)) {
                pdDocument.save(out, CompressParameters.NO_COMPRESSION);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        });
    }


    private static void createPageResourcesIfNotPresent(PDPage pdPage) {

        PDResources resources = pdPage.getResources();
        if (resources == null) {
            resources = new PDResources();
            pdPage.setResources(resources);
        }
    }


    private static AffineTransform getTextDeRotationTransform(PDPage page) {

        return AffineTransform.getQuadrantRotateInstance(switch (page.getRotation()) {
            case 90 -> 3;
            case 180 -> 2;
            case 270 -> 1;
            default -> 0;
        });
    }

}
