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

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.File;
import java.nio.file.Path;
import java.util.Collection;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.function.Predicate;

import org.apache.commons.lang3.StringUtils;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;

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.LayoutParsingPipeline;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationDocument;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Document;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Section;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.SemanticNode;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.SuperSection;
import com.knecon.fforesight.service.layoutparser.processor.model.outline.OutlineObject;
import com.knecon.fforesight.service.layoutparser.processor.model.outline.OutlineObjectTree;
import com.knecon.fforesight.service.layoutparser.processor.model.outline.TableOfContents;
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.factory.DocumentGraphFactory;
import com.knecon.fforesight.service.layoutparser.processor.services.visualization.LayoutGridService;
import com.knecon.fforesight.service.layoutparser.server.utils.AbstractTest;
import com.knecon.fforesight.service.viewerdoc.service.PDFTronViewerDocumentService;

import jakarta.annotation.PostConstruct;
import lombok.SneakyThrows;

public class OutlineDetectionTest extends AbstractTest {

    PDFTronViewerDocumentService viewerDocumentService = new PDFTronViewerDocumentService(null);
    LayoutGridService layoutGridService = new LayoutGridService(viewerDocumentService);

    @Autowired
    protected LayoutParsingPipeline layoutParsingPipeline;

    @Autowired
    PDFNetInitializer pdfNetInitializer;


    @BeforeEach
    public void init() {

        pdfNetInitializer.init();
    }


    @Test
    @SneakyThrows
    public void testOutlinesToSections() {

        String fileName = "files/new/crafted_outline_test_doc.pdf";

        String tmpFileName = "/tmp/" + Path.of(fileName).getFileName() + "_VIEWER.pdf";

        var documentFile = new ClassPathResource(fileName).getFile();

        long start = System.currentTimeMillis();
        ClassificationDocument classificationDocument = parseLayout(fileName, LayoutParsingType.REDACT_MANAGER_WITHOUT_DUPLICATE_PARAGRAPH);

        OutlineObjectTree outlineObjectTree = classificationDocument.getOutlineObjectTree();
        assertEquals(outlineObjectTree.getRootNodes().size(), 8);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(1).size(), 1);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(3).size(), 1);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(4).size(), 1);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(5).size(), 1);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(6).size(), 2);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(7).size(), 3);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(8).size(), 2);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(10).size(), 1);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(11).size(), 4);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(12).size(), 1);
        assertEquals(outlineObjectTree.getOutlineObjectsPerPage().get(13).size(), 2);
        assertTrue(outlineObjectTree.getOutlineObjectsPerPage().values()
                           .stream()
                           .flatMap(Collection::stream)
                           .allMatch(OutlineObject::isFound));

        TableOfContents tableOfContents = classificationDocument.getTableOfContents();

        assertEquals(tableOfContents.getMainSections().size(), 9);
        assertEquals(tableOfContents.getMainSections().subList(1, 9)
                             .stream()
                             .map(tableOfContentItem -> sanitizeString(tableOfContentItem.getHeadline().toString()))
                             .toList(),
                     outlineObjectTree.getRootNodes()
                             .stream()
                             .map(outlineObjectTreeNode -> sanitizeString(outlineObjectTreeNode.getOutlineObject().getTitle()))
                             .toList());
        assertEquals(tableOfContents.getMainSections().get(5).getChildren().size(), 6);
        assertEquals(tableOfContents.getMainSections().get(7).getChildren().size(), 3);
        assertEquals(tableOfContents.getMainSections().get(8).getChildren().size(), 3);
        assertEquals(tableOfContents.getMainSections().get(8).getChildren().get(2).getChildren().size(), 1);
        assertEquals(tableOfContents.getMainSections().get(8).getChildren().get(2).getChildren().get(0).getChildren().size(), 3);

        assertEquals(tableOfContents.getMainSections().get(0).getImages().size(), 1);
        assertEquals(tableOfContents.getMainSections().get(6).getImages().size(), 1);
        assertEquals(tableOfContents.getMainSections().get(8).getChildren().get(2).getChildren().get(0).getChildren().get(2).getImages().size(), 1);

        Document document = buildGraph(fileName, classificationDocument);

        assertTrue(tableOfContents.getAllTableOfContentItems()
                           .stream()
                           .allMatch(tableOfContentItem -> tableOfContentItem.getSection() != null));
        assertTrue(tableOfContents.getAllTableOfContentItems()
                           .stream()
                           .filter(tableOfContentItem -> tableOfContentItem.getChildren().isEmpty())
                           .allMatch(tableOfContentItem -> tableOfContentItem.getSection() instanceof Section));
        assertTrue(tableOfContents.getAllTableOfContentItems()
                           .stream()
                           .filter(tableOfContentItem -> !tableOfContentItem.getChildren().isEmpty())
                           .allMatch(tableOfContentItem -> tableOfContentItem.getSection() instanceof SuperSection));

        List<SemanticNode> childrenOfTypeSectionOrSuperSection = document.getChildrenOfTypeSectionOrSuperSection();

        assertEquals(childrenOfTypeSectionOrSuperSection.size(), 9);
        assertEquals(childrenOfTypeSectionOrSuperSection.subList(1, 9)
                             .stream()
                             .map(section -> sanitizeString(section.getHeadline().getLeafTextBlock().toString()))
                             .toList(),
                     outlineObjectTree.getRootNodes()
                             .stream()
                             .map(outlineObjectTreeNode -> sanitizeString(outlineObjectTreeNode.getOutlineObject().getTitle()))
                             .toList());
        Predicate<SemanticNode> isSectionOrSuperSection = semanticNode -> semanticNode instanceof Section || semanticNode instanceof SuperSection;
        assertEquals(childrenOfTypeSectionOrSuperSection.get(5).streamChildren()
                             .filter(isSectionOrSuperSection)
                             .count(), 6 + 1); // 1 additional for main text of parent section
        assertEquals(childrenOfTypeSectionOrSuperSection.get(7).streamChildren()
                             .filter(isSectionOrSuperSection)
                             .count(), 3 + 1);
        assertEquals(childrenOfTypeSectionOrSuperSection.get(8).streamChildren()
                             .filter(isSectionOrSuperSection)
                             .count(), 3 + 1);
        assertEquals(childrenOfTypeSectionOrSuperSection.get(8).streamChildren()
                             .filter(isSectionOrSuperSection)
                             .toList().get(3).streamChildren()
                             .filter(isSectionOrSuperSection)
                             .count(), 1 + 1);
        assertEquals(childrenOfTypeSectionOrSuperSection.get(8).streamChildren()
                             .filter(isSectionOrSuperSection)
                             .toList().get(3).streamChildren()
                             .filter(isSectionOrSuperSection)
                             .toList().get(1).streamChildren()
                             .filter(isSectionOrSuperSection)
                             .count(), 3 + 1);

        List<List<Integer>> imageTreeIdList = document.streamAllImages()
                .map(image -> image.getParent().getTreeId())
                .toList();

        assertEquals(imageTreeIdList.get(0), List.of(0));
        assertEquals(imageTreeIdList.get(1), List.of(6));
        assertEquals(imageTreeIdList.get(2), List.of(8, 4, 2, 4));

        layoutGridService.addLayoutGrid(documentFile, document, new File(tmpFileName), true, false);
        System.out.printf("Total time: %.2fs%n", ((float) (System.currentTimeMillis() - start)) / 1000);
    }


    private static String sanitizeString(String text) {

        return StringUtils.deleteWhitespace(text).toLowerCase(Locale.ROOT);
    }


    @SneakyThrows
    protected ClassificationDocument parseLayout(String filename, LayoutParsingType layoutParsingType) {

        File fileResource = new ClassPathResource(filename).getFile();
        prepareStorage(filename, "cv_table_parsing_response/empty.json", "image_service_response/crafted_outline_test_doc.IMAGE_INFO.json");
        return layoutParsingPipeline.parseLayout(layoutParsingType,
                                                 fileResource,
                                                 layoutParsingStorageService.getImagesFile(IMAGE_FILE_ID),
                                                 new TableServiceResponse(),
                                                 new VisualLayoutParsingResponse(),
                                                 Map.of("file", filename, "debug", "true"));
    }


    @SneakyThrows
    protected Document buildGraph(String filename, ClassificationDocument classificationDocument) {

        if (!filename.startsWith("files") && filename.startsWith("/")) {
            LayoutParsingRequest layoutParsingRequest = buildDefaultLayoutParsingRequest(Path.of(filename).getFileName().toString(), LayoutParsingType.REDACT_MANAGER, true);
            prepareStorage(layoutParsingRequest, new File(filename));
            return DocumentGraphFactory.buildDocumentGraph(LayoutParsingType.REDACT_MANAGER,
                                                           layoutParsingPipeline.parseLayout(LayoutParsingType.REDACT_MANAGER,
                                                                                             new File(filename),
                                                                                             layoutParsingStorageService.getImagesFile(layoutParsingRequest.imagesFileStorageId()
                                                                                                                                               .get()),
                                                                                             new TableServiceResponse(),
                                                                                             new VisualLayoutParsingResponse(),
                                                                                             layoutParsingRequest.identifier()));
        } else {
            prepareStorage(filename);
            return DocumentGraphFactory.buildDocumentGraph(LayoutParsingType.REDACT_MANAGER, classificationDocument);
        }

    }

}

