package com.knecon.fforesight.service.layoutparser.processor;

import static java.lang.String.format;

import java.awt.geom.Rectangle2D;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

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.apache.pdfbox.pdmodel.documentinterchange.markedcontent.PDMarkedContent;
import org.springframework.stereotype.Service;

import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.NodeType;
import com.knecon.fforesight.service.layoutparser.internal.api.queue.LayoutParsingFinishedEvent;
import com.knecon.fforesight.service.layoutparser.internal.api.queue.LayoutParsingRequest;
import com.knecon.fforesight.service.layoutparser.internal.api.queue.LayoutParsingType;
import com.knecon.fforesight.service.layoutparser.processor.model.AbstractPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationDocument;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationPage;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Document;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.ImageType;
import com.knecon.fforesight.service.layoutparser.processor.model.image.ClassifiedImage;
import com.knecon.fforesight.service.layoutparser.processor.model.table.Cell;
import com.knecon.fforesight.service.layoutparser.processor.model.table.CleanRulings;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPositionSequence;
import com.knecon.fforesight.service.layoutparser.processor.python_api.adapter.CvTableParsingAdapter;
import com.knecon.fforesight.service.layoutparser.processor.python_api.adapter.ImageServiceResponseAdapter;
import com.knecon.fforesight.service.layoutparser.processor.python_api.adapter.VisualLayoutParsingAdapter;
import com.knecon.fforesight.service.layoutparser.processor.python_api.model.image.ImageServiceResponse;
import com.knecon.fforesight.service.layoutparser.processor.python_api.model.table.TableCells;
import com.knecon.fforesight.service.layoutparser.processor.python_api.model.table.TableServiceResponse;
import com.knecon.fforesight.service.layoutparser.processor.python_api.model.table.VisualLayoutParsingResponse;
import com.knecon.fforesight.service.layoutparser.processor.services.BodyTextFrameService;
import com.knecon.fforesight.service.layoutparser.processor.services.RulingCleaningService;
import com.knecon.fforesight.service.layoutparser.processor.services.SectionsBuilderService;
import com.knecon.fforesight.service.layoutparser.processor.services.SimplifiedSectionTextService;
import com.knecon.fforesight.service.layoutparser.processor.services.TableExtractionService;
import com.knecon.fforesight.service.layoutparser.processor.services.TextRulingsClassifier;
import com.knecon.fforesight.service.layoutparser.processor.services.blockification.DocstrumBlockificationService;
import com.knecon.fforesight.service.layoutparser.processor.services.blockification.DocuMineBlockificationService;
import com.knecon.fforesight.service.layoutparser.processor.services.blockification.RedactManagerBlockificationService;
import com.knecon.fforesight.service.layoutparser.processor.services.classification.ClarifyndClassificationService;
import com.knecon.fforesight.service.layoutparser.processor.services.classification.DocuMineClassificationService;
import com.knecon.fforesight.service.layoutparser.processor.services.classification.RedactManagerClassificationService;
import com.knecon.fforesight.service.layoutparser.processor.services.factory.DocumentGraphFactory;
import com.knecon.fforesight.service.layoutparser.processor.services.graphics.GraphicExtractorService;
import com.knecon.fforesight.service.layoutparser.processor.services.mapper.DocumentDataMapper;
import com.knecon.fforesight.service.layoutparser.processor.services.mapper.TaasDocumentDataMapper;
import com.knecon.fforesight.service.layoutparser.processor.services.parsing.PDFLinesTextStripper;
import com.knecon.fforesight.service.layoutparser.processor.services.visualization.LayoutGridService;
import com.knecon.fforesight.service.layoutparser.processor.utils.MarkedContentUtils;

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

@SuppressWarnings("PMD.CloseResource")
@Slf4j
@Service
@RequiredArgsConstructor
@FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
public class LayoutParsingPipeline {

    ImageServiceResponseAdapter imageServiceResponseAdapter;
    CvTableParsingAdapter cvTableParsingAdapter;
    LayoutParsingStorageService layoutParsingStorageService;
    SectionsBuilderService sectionsBuilderService;
    RedactManagerClassificationService redactManagerClassificationService;
    DocuMineClassificationService docuMineClassificationService;
    SimplifiedSectionTextService simplifiedSectionTextService;
    BodyTextFrameService bodyTextFrameService;
    RulingCleaningService rulingCleaningService;
    TableExtractionService tableExtractionService;
    DocuMineBlockificationService docuMineBlockificationService;
    RedactManagerBlockificationService redactManagerBlockificationService;
    DocstrumBlockificationService docstrumBlockificationService;
    LayoutGridService layoutGridService;
    ObservationRegistry observationRegistry;
    VisualLayoutParsingAdapter visualLayoutParsingAdapter;
    ClarifyndClassificationService clarifyndClassificationService;
    GraphicExtractorService graphicExtractorService;


    public LayoutParsingFinishedEvent parseLayoutAndSaveFilesToStorage(LayoutParsingRequest layoutParsingRequest) throws IOException {

        long start = System.currentTimeMillis();
        log.info("Starting layout parsing for {}", layoutParsingRequest.identifier());

        File originFile = layoutParsingStorageService.getOriginFile(layoutParsingRequest.originFileStorageId());
        File viewerDocumentFile = layoutParsingStorageService.getViewerDocFile(layoutParsingRequest.viewerDocumentStorageId()).orElse(originFile);

        VisualLayoutParsingResponse visualLayoutParsingResponse = new VisualLayoutParsingResponse();
        if (layoutParsingRequest.visualLayoutParsingFileId().isPresent()) {
            visualLayoutParsingResponse = layoutParsingStorageService.getVisualLayoutParsingFile(layoutParsingRequest.visualLayoutParsingFileId().get());
        }

        ImageServiceResponse imageServiceResponse = new ImageServiceResponse();
        if (layoutParsingRequest.imagesFileStorageId().isPresent()) {
            imageServiceResponse = layoutParsingStorageService.getImagesFile(layoutParsingRequest.imagesFileStorageId().get());
        }

        TableServiceResponse tableServiceResponse = new TableServiceResponse();
        if (layoutParsingRequest.tablesFileStorageId().isPresent()) {
            tableServiceResponse = layoutParsingStorageService.getTablesFile(layoutParsingRequest.tablesFileStorageId().get());
        }

        ClassificationDocument classificationDocument = parseLayout(layoutParsingRequest.layoutParsingType(),
                                                                    originFile,
                                                                    imageServiceResponse,
                                                                    tableServiceResponse,
                                                                    visualLayoutParsingResponse,
                                                                    layoutParsingRequest.identifier());

        log.info("Building document graph for {}", layoutParsingRequest.identifier());

        Document documentGraph = observeBuildDocumentGraph(layoutParsingRequest.layoutParsingType(), classificationDocument);

        log.info("Creating viewer document for {}", layoutParsingRequest.identifier());

        layoutGridService.addLayoutGrid(viewerDocumentFile, documentGraph, viewerDocumentFile, false, layoutParsingRequest.visualLayoutParsingFileId().isPresent());

        log.info("Storing resulting files for {}", layoutParsingRequest.identifier());

        layoutParsingStorageService.storeDocumentData(layoutParsingRequest, DocumentDataMapper.toDocumentData(documentGraph));
        layoutParsingStorageService.storeSimplifiedText(layoutParsingRequest, simplifiedSectionTextService.toSimplifiedText(documentGraph));
        layoutParsingStorageService.storeViewerDocument(layoutParsingRequest, viewerDocumentFile);

        if (layoutParsingRequest.layoutParsingType().equals(LayoutParsingType.CLARIFYND)) {
            log.info("Building research document data for {}", layoutParsingRequest.identifier());
            var researchDocumentData = TaasDocumentDataMapper.fromDocument(documentGraph);
            layoutParsingStorageService.storeResearchDocumentData(layoutParsingRequest, researchDocumentData);
        }

        if (!viewerDocumentFile.equals(originFile)) {
            viewerDocumentFile.delete();
        }
        originFile.delete();

        return LayoutParsingFinishedEvent.builder()
                .identifier(layoutParsingRequest.identifier())
                .numberOfPages(documentGraph.getNumberOfPages())
                .duration(System.currentTimeMillis() - start)
                .message(format("""
                                        Layout parsing has finished in %.02f s.
                                        identifiers: %s
                                        %s
                                        Files have been saved with Ids:
                                        Structure: %s
                                        Text: %s
                                        Positions: %s
                                        PageData: %s
                                        Simplified Text: %s
                                        Viewer Doc: %s""",
                                ((float) (System.currentTimeMillis() - start)) / 1000,
                                layoutParsingRequest.identifier(),
                                buildSemanticNodeCountMessage(documentGraph.getNumberOfPages(), documentGraph.buildSemanticNodeCounts()),
                                layoutParsingRequest.structureFileStorageId(),
                                layoutParsingRequest.textBlockFileStorageId(),
                                layoutParsingRequest.positionBlockFileStorageId(),
                                layoutParsingRequest.pageFileStorageId(),
                                layoutParsingRequest.simplifiedTextStorageId(),
                                layoutParsingRequest.viewerDocumentStorageId()))
                .build();

    }


    private Document observeBuildDocumentGraph(LayoutParsingType layoutParsingType, ClassificationDocument classificationDocument) {

        AtomicReference<Document> documentReference = new AtomicReference<>();

        Observation.createNotStarted("LayoutParsingPipeline", observationRegistry)
                .contextualName("build-document-graph")
                .observe(() -> documentReference.set(DocumentGraphFactory.buildDocumentGraph(layoutParsingType, classificationDocument)));

        return documentReference.get();
    }


    private String buildSemanticNodeCountMessage(int numberOfPages, Map<NodeType, Long> semanticNodeCounts) {

        return String.format("%d pages with %d sections, %d headlines, %d paragraphs, %d tables with %d cells, %d headers, and %d footers parsed",
                             numberOfPages,
                             semanticNodeCounts.get(NodeType.SECTION) == null ? 0 : semanticNodeCounts.get(NodeType.SECTION),
                             semanticNodeCounts.get(NodeType.HEADLINE) == null ? 0 : semanticNodeCounts.get(NodeType.HEADLINE),
                             semanticNodeCounts.get(NodeType.PARAGRAPH) == null ? 0 : semanticNodeCounts.get(NodeType.PARAGRAPH),
                             semanticNodeCounts.get(NodeType.TABLE) == null ? 0 : semanticNodeCounts.get(NodeType.TABLE),
                             semanticNodeCounts.get(NodeType.TABLE_CELL) == null ? 0 : semanticNodeCounts.get(NodeType.TABLE_CELL),
                             semanticNodeCounts.get(NodeType.HEADER) == null ? 0 : semanticNodeCounts.get(NodeType.HEADER),
                             semanticNodeCounts.get(NodeType.FOOTER) == null ? 0 : semanticNodeCounts.get(NodeType.FOOTER));
    }


    @SneakyThrows
    @Observed(name = "LayoutParsingPipeline", contextualName = "parse-layout")
    public ClassificationDocument parseLayout(LayoutParsingType layoutParsingType,
                                              File originFile,
                                              ImageServiceResponse imageServiceResponse,
                                              TableServiceResponse tableServiceResponse,
                                              VisualLayoutParsingResponse visualLayoutParsingResponse,
                                              Map<String, String> identifier) {

        PDDocument originDocument = openDocument(originFile);
        addNumberOfPagesToTrace(originDocument.getNumberOfPages(), Files.size(originFile.toPath()));
        Map<Integer, List<TableCells>> pdfTableCells = cvTableParsingAdapter.buildCvParsedTablesPerPage(tableServiceResponse);
        Map<Integer, List<ClassifiedImage>> pdfImages = imageServiceResponseAdapter.buildClassifiedImagesPerPage(imageServiceResponse);
        Map<Integer, List<ClassifiedImage>> signatures = visualLayoutParsingAdapter.buildExtractedSignaturesPerPage(visualLayoutParsingResponse);
        ClassificationDocument classificationDocument = new ClassificationDocument();

        classificationDocument.getVisualizations().setActive(identifier.containsKey("debug"));

        List<ClassificationPage> classificationPages = new ArrayList<>();

        long pageCount = originDocument.getNumberOfPages();

        for (int pageNumber = 1; pageNumber <= pageCount; pageNumber++) {

            if (pageNumber % 100 == 0) {
                // re-open document every once in a while to save on RAM. This has no significant performance impact.
                // This is due to PDFBox caching all images and some other stuff with Soft References. This dereferences them and forces the freeing of memory.
                originDocument.close();
                originDocument = openDocument(originFile);
            }

            if (pageNumber % 100 == 0 || pageNumber == pageCount || pageNumber == 1) {
                log.info("Extracting text on Page {} for {}", pageNumber, identifier);
            }

            classificationDocument.setPages(classificationPages);
            PDFLinesTextStripper stripper = new PDFLinesTextStripper();
            PDPage pdPage = originDocument.getPage(pageNumber - 1);
            stripper.setPageNumber(pageNumber);
            stripper.setStartPage(pageNumber);
            stripper.setEndPage(pageNumber);
            stripper.setPdpage(pdPage);
            if (layoutParsingType.equals(LayoutParsingType.DOCUMINE_OLD)) {
                stripper.setSortByPosition(true);
            }
            stripper.getText(originDocument);
            List<TextPositionSequence> words = stripper.getTextPositionSequences();
            classificationDocument.getVisualizations().addTextVisualizations(words, pageNumber);

            PDRectangle pdr = pdPage.getMediaBox();

            int rotation = pdPage.getRotation();
            boolean isLandscape = pdr.getWidth() > pdr.getHeight() && (rotation == 0 || rotation == 180) || pdr.getHeight() > pdr.getWidth() && (rotation == 90 || rotation == 270);

            PDRectangle cropbox = pdPage.getCropBox();
            CleanRulings cleanRulings = rulingCleaningService.getCleanRulings(pdfTableCells.get(pageNumber), stripper.getRulings());

            List<Cell> emptyTableCells = TableExtractionService.findCells(cleanRulings.getHorizontals(), cleanRulings.getVerticals());
            classificationDocument.getVisualizations().addCellVisualizations(emptyTableCells, pageNumber);

            TextRulingsClassifier.classifyUnderlinedAndStrikethroughText(words, cleanRulings);

            var graphics = graphicExtractorService.extractPathElementGraphics(originDocument, pdPage, pageNumber, cleanRulings, stripper.getTextPositionSequences(), false);

            pdfImages.computeIfAbsent(pageNumber, x -> new ArrayList<>())
                    .addAll(graphics.stream()
                                    .map(g -> new ClassifiedImage(new Rectangle2D.Double(g.x1, g.y1, g.width(), g.height()), ImageType.GRAPHIC, false, stripper.getPageNumber()))
                                    .toList());

            ClassificationPage classificationPage = switch (layoutParsingType) {
                case REDACT_MANAGER_OLD ->
                        redactManagerBlockificationService.blockify(stripper.getTextPositionSequences(), cleanRulings, classificationDocument.getVisualizations());
                case DOCUMINE_OLD -> docuMineBlockificationService.blockify(words, cleanRulings);
                case DOCUMINE, REDACT_MANAGER, REDACT_MANAGER_PARAGRAPH_DEBUG ->
                        docstrumBlockificationService.blockify(words, cleanRulings, true, classificationDocument.getVisualizations());
                case CLARIFYND, CLARIFYND_PARAGRAPH_DEBUG -> docstrumBlockificationService.blockify(words, cleanRulings, false, classificationDocument.getVisualizations());
            };

            classificationPage.setCleanRulings(cleanRulings);
            classificationPage.setRotation(rotation);
            classificationPage.setLandscape(isLandscape);
            classificationPage.setPageNumber(pageNumber);
            classificationPage.setPageWidth(cropbox.getWidth());
            classificationPage.setPageHeight(cropbox.getHeight());

            classificationDocument.getVisualizations().addMarkedContentVisualizations(stripper.getMarkedContents(), pageNumber, pdPage);
            // MarkedContent needs to be converted at this point, otherwise it leads to GC Problems in Pdfbox.
            classificationPage.setMarkedContentBboxPerType(convertMarkedContents(stripper.getMarkedContents(), pdPage));

            // If images is ocr needs to be calculated before textBlocks are moved into tables, otherwise findOcr algorithm needs to be adopted.
            if (pdfImages != null && pdfImages.containsKey(pageNumber)) {
                classificationPage.setImages(pdfImages.get(pageNumber));
                imageServiceResponseAdapter.findOcr(classificationPage);
            }

            if (signatures.containsKey(pageNumber)) {
                if (classificationPage.getImages() == null || classificationPage.getImages().isEmpty()) {
                    classificationPage.setImages(signatures.get(pageNumber));
                } else {
                    classificationPage.getImages().addAll(signatures.get(pageNumber));
                }
            }

            tableExtractionService.extractTables(emptyTableCells, classificationPage);

            if (layoutParsingType == LayoutParsingType.REDACT_MANAGER) {
                docstrumBlockificationService.combineBlocks(classificationPage);
            } else if (layoutParsingType == LayoutParsingType.CLARIFYND) {
                docstrumBlockificationService.mergeIntersectingBlocks(classificationPage, 0, 6.5f);
            }

            buildPageStatistics(classificationPage);
            increaseDocumentStatistics(classificationPage, classificationDocument);

            classificationPages.add(classificationPage);
        }

        originDocument.close();

        log.info("Calculating BodyTextFrame for {}", identifier);
        bodyTextFrameService.setBodyTextFrames(classificationDocument, layoutParsingType);
        for (ClassificationPage page : classificationDocument.getPages()) {
            classificationDocument.getVisualizations().addCleanRulingVisualization(page.getCleanRulings(), page.getPageNumber());
        }
        log.info("Classify TextBlocks for {}", identifier);
        switch (layoutParsingType) {
            case REDACT_MANAGER, REDACT_MANAGER_PARAGRAPH_DEBUG, REDACT_MANAGER_OLD, CLARIFYND_PARAGRAPH_DEBUG ->
                    redactManagerClassificationService.classifyDocument(classificationDocument);
            case DOCUMINE_OLD, DOCUMINE -> docuMineClassificationService.classifyDocument(classificationDocument);
            case CLARIFYND -> clarifyndClassificationService.classifyDocument(classificationDocument);
        }

        log.info("Building Sections for {}", identifier);

        switch (layoutParsingType) {
            case CLARIFYND_PARAGRAPH_DEBUG, REDACT_MANAGER_PARAGRAPH_DEBUG -> sectionsBuilderService.buildParagraphDebugSections(classificationDocument);
            default -> {
                sectionsBuilderService.buildSections(classificationDocument);
                sectionsBuilderService.addImagesToSections(classificationDocument);
            }
        }

        return classificationDocument;
    }


    private void addNumberOfPagesToTrace(int numberOfPages, long size) {

        if (observationRegistry.getCurrentObservation() != null) {
            observationRegistry.getCurrentObservation().highCardinalityKeyValue("numberOfPages", String.valueOf(numberOfPages));
            observationRegistry.getCurrentObservation().highCardinalityKeyValue("fileSize", String.valueOf(size));
        }
    }


    @SneakyThrows
    private PDDocument openDocument(File originFile) {

        PDDocument document = Loader.loadPDF(originFile);
        document.setAllSecurityToBeRemoved(true);
        return document;
    }


    private Map<String, List<Rectangle2D>> convertMarkedContents(List<PDMarkedContent> pdMarkedContents, PDPage pdPage) {

        Map<String, List<Rectangle2D>> markedContentBboxes = new HashMap<>();
        markedContentBboxes.put(MarkedContentUtils.HEADER, MarkedContentUtils.getMarkedContentBboxPerLine(pdMarkedContents, MarkedContentUtils.HEADER, pdPage));
        markedContentBboxes.put(MarkedContentUtils.FOOTER, MarkedContentUtils.getMarkedContentBboxPerLine(pdMarkedContents, MarkedContentUtils.FOOTER, pdPage));
        return markedContentBboxes;
    }


    private void increaseDocumentStatistics(ClassificationPage classificationPage, ClassificationDocument document) {

        document.getFontSizeCounter().addAll(classificationPage.getFontSizeCounter().getCountPerValue());
        document.getFontCounter().addAll(classificationPage.getFontCounter().getCountPerValue());
        document.getTextHeightCounter().addAll(classificationPage.getTextHeightCounter().getCountPerValue());
        document.getFontStyleCounter().addAll(classificationPage.getFontStyleCounter().getCountPerValue());
    }


    private void buildPageStatistics(ClassificationPage classificationPage) {

        // Collect all statistics for the classificationPage, except from blocks inside tables, as tables will always be added to BodyTextFrame.
        for (AbstractPageBlock textBlock : classificationPage.getTextBlocks()) {
            if (textBlock instanceof TextPageBlock) {
                if (((TextPageBlock) textBlock).getSequences() == null) {
                    continue;
                }
                for (TextPositionSequence word : ((TextPageBlock) textBlock).getSequences()) {
                    classificationPage.getTextHeightCounter().add(word.getTextHeight());
                    classificationPage.getFontCounter().add(word.getFont());
                    classificationPage.getFontSizeCounter().add(word.getFontSize());
                    classificationPage.getFontStyleCounter().add(word.getFontStyle());
                }
            }
        }

    }

}
