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

import static java.lang.String.format;
import static java.util.stream.Collectors.groupingBy;
import static java.util.stream.Collectors.toList;

import java.awt.geom.Rectangle2D;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.stream.Collectors;

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.ClassificationFooter;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationHeader;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationPage;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.DocumentTree;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Document;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.DuplicatedParagraph;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Footer;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.GenericSemanticNode;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Header;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Headline;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Image;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Page;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Paragraph;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Section;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.textblock.AtomicTextBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.image.ClassifiedImage;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.utils.IdBuilder;
import com.knecon.fforesight.service.layoutparser.processor.utils.TextPositionOperations;

import lombok.AccessLevel;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Getter;
import lombok.experimental.FieldDefaults;
import lombok.experimental.UtilityClass;

@UtilityClass
public class DocumentGraphFactory {

    public Document buildDocumentGraph(LayoutParsingType layoutParsingType, ClassificationDocument document) {

        Document documentGraph = new Document();

        documentGraph.setVisualizations(document.getVisualizations());

        Context context = new Context(documentGraph);

        document.getPages()
                .forEach(context::buildAndAddPageWithCounter);
        document.getSections()
                .stream()
                .flatMap(section -> section.getImages()
                        .stream())
                .forEach(image -> context.getImages().add(image));
        addSections(layoutParsingType, document, context, documentGraph);
        addHeaderAndFooterToEachPage(document, context);

        documentGraph.setNumberOfPages(context.pages.size());
        documentGraph.setPages(context.pages.keySet());
        documentGraph.setDocumentTree(context.documentTree);
        documentGraph.setTextBlock(documentGraph.getTextBlock());
        return documentGraph;
    }


    private void addSections(LayoutParsingType layoutParsingType, ClassificationDocument classificationDocument, Context context, Document document) {

        classificationDocument.getSections()
                .forEach(section -> SectionNodeFactory.addSection(layoutParsingType, null, section.getNonEmptyPageBlocks(), section.getImages(), context, document));
    }


    public void addParagraphOrHeadline(GenericSemanticNode parentNode,
                                       TextPageBlock originalTextBlock,
                                       Context context,
                                       List<TextPageBlock> textBlocksToMerge,
                                       LayoutParsingType layoutParsingType) {

        Page page = context.getPage(originalTextBlock.getPage());

        GenericSemanticNode node;
        if (originalTextBlock.isHeadline()) {
            node = Headline.builder().documentTree(context.getDocumentTree()).build();
        } else if (originalTextBlock.isToDuplicate() && layoutParsingType.equals(LayoutParsingType.REDACT_MANAGER)) {
            node = DuplicatedParagraph.builder().documentTree(context.getDocumentTree()).build();
        } else {
            node = Paragraph.builder().documentTree(context.getDocumentTree()).build();
        }

        page.getMainBody().add(node);

        List<TextPageBlock> textBlocks = new ArrayList<>();
        textBlocks.add(originalTextBlock);
        textBlocks.addAll(textBlocksToMerge);

        AtomicTextBlock textBlock = context.textBlockFactory.buildAtomicTextBlock(TextPositionOperations.mergeAndSortTextPositionSequenceByYThenX(textBlocks), node, context, page);

        if (node instanceof DuplicatedParagraph duplicatedParagraph) {
            AtomicTextBlock unsortedTextBlock = context.textBlockFactory.buildAtomicTextBlock(textBlocks.stream()
                                                                                                      .flatMap(tb -> tb.getSequences()
                                                                                                              .stream())
                                                                                                      .collect(Collectors.toList()), node, context, page);
            duplicatedParagraph.setUnsortedLeafTextBlock(unsortedTextBlock);
        }

        List<Integer> treeId = context.documentTree.createNewChildEntryAndReturnId(parentNode, node);
        node.setLeafTextBlock(textBlock);
        node.setTreeId(treeId);
    }


    public void addImage(GenericSemanticNode parent, ClassifiedImage image, Context context) {

        Image imageNode = createImage(image, context);
        List<Integer> treeId = context.getDocumentTree().createNewChildEntryAndReturnId(parent, imageNode);
        imageNode.setTreeId(treeId);
        imageNode.setLeafTextBlock(context.textBlockFactory.emptyTextBlock(parent, context, context.getPage(image.getPage())));
    }


    private Image createImage(ClassifiedImage image, Context context) {

        Rectangle2D position = image.getPosition();
        Page page = context.getPage(image.getPage());
        Image imageNode = Image.builder()
                .id(IdBuilder.buildId(Set.of(page), List.of(position)))
                .imageType(image.getImageType())
                .position(position)
                .transparent(image.isHasTransparency())
                .page(page)
                .documentTree(context.getDocumentTree())
                .build();
        page.getMainBody().add(imageNode);
        return imageNode;
    }


    private void addHeaderAndFooterToEachPage(ClassificationDocument document, Context context) {

        Map<Integer, List<TextPageBlock>> headers = document.getHeaders()
                .stream()
                .map(ClassificationHeader::getTextBlocks)
                .flatMap(List::stream)
                .collect(groupingBy(AbstractPageBlock::getPage, toList()));

        Map<Integer, List<TextPageBlock>> footers = document.getFooters()
                .stream()
                .map(ClassificationFooter::getTextBlocks)
                .flatMap(List::stream)
                .collect(groupingBy(AbstractPageBlock::getPage, toList()));

        for (int pageIndex = 1; pageIndex <= document.getPages().size(); pageIndex++) {
            if (headers.containsKey(pageIndex)) {
                addHeader(headers.get(pageIndex), context);
            } else {
                addEmptyHeader(pageIndex, context);
            }
        }

        for (int pageIndex = 1; pageIndex <= document.getPages().size(); pageIndex++) {
            if (footers.containsKey(pageIndex)) {
                addFooter(footers.get(pageIndex), context);
            } else {
                addEmptyFooter(pageIndex, context);
            }
        }
    }


    private void addFooter(List<TextPageBlock> textBlocks, Context context) {

        Page page = context.getPage(textBlocks.get(0).getPage());
        Footer footer = Footer.builder().documentTree(context.getDocumentTree()).build();
        AtomicTextBlock textBlock = context.textBlockFactory.buildAtomicTextBlock(TextPositionOperations.mergeTextPositionSequence(textBlocks),
                                                                                  footer,
                                                                                  context,
                                                                                  page);
        List<Integer> tocId = context.getDocumentTree().createNewMainEntryAndReturnId(footer);
        footer.setTreeId(tocId);
        footer.setLeafTextBlock(textBlock);
        page.setFooter(footer);
    }


    public void addHeader(List<TextPageBlock> textBlocks, Context context) {

        Page page = context.getPage(textBlocks.get(0).getPage());
        Header header = Header.builder().documentTree(context.getDocumentTree()).build();
        AtomicTextBlock textBlock = context.textBlockFactory.buildAtomicTextBlock(TextPositionOperations.mergeTextPositionSequence(textBlocks), header, 0, page);
        List<Integer> tocId = context.getDocumentTree().createNewMainEntryAndReturnId(header);
        header.setTreeId(tocId);
        header.setLeafTextBlock(textBlock);
        page.setHeader(header);
    }


    private void addEmptyFooter(int pageIndex, Context context) {

        Page page = context.getPage(pageIndex);
        Footer footer = Footer.builder().documentTree(context.getDocumentTree()).build();
        AtomicTextBlock textBlock = context.textBlockFactory.emptyTextBlock(footer, context, page);
        List<Integer> tocId = context.getDocumentTree().createNewMainEntryAndReturnId(footer);
        footer.setTreeId(tocId);
        footer.setLeafTextBlock(textBlock);
        page.setFooter(footer);
    }


    private void addEmptyHeader(int pageIndex, Context context) {

        Page page = context.getPage(pageIndex);
        Header header = Header.builder().documentTree(context.getDocumentTree()).build();
        AtomicTextBlock textBlock = context.textBlockFactory.emptyTextBlock(header, 0, page);
        List<Integer> tocId = context.getDocumentTree().createNewMainEntryAndReturnId(header);
        header.setTreeId(tocId);
        header.setLeafTextBlock(textBlock);
        page.setHeader(header);
    }


    @Getter
    @Builder
    @AllArgsConstructor
    @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
    public final class Context {

        DocumentTree documentTree;
        Map<Page, Integer> pages;
        List<Section> sections;
        List<ClassifiedImage> images;
        TextBlockFactory textBlockFactory;


        public Context(Document document) {

            documentTree = new DocumentTree(document);
            pages = new HashMap<>();
            sections = new LinkedList<>();
            images = new LinkedList<>();
            textBlockFactory = new TextBlockFactory();
        }


        public void buildAndAddPageWithCounter(ClassificationPage classificationPage) {

            Page page = Page.fromClassificationPage(classificationPage);
            //this counter counts the TextBlocks per page
            //initial value is set to 1, because 0 is reserved for Header
            pages.put(page, 1);
        }


        public int getAndIncrementTextBlockNumberOnPage(Page page) {

            Integer textBlockNumberOnPage = pages.get(page);
            pages.merge(page, 1, Integer::sum);
            return textBlockNumberOnPage;
        }


        public Page getPage(int pageIndex) {

            return pages.keySet()
                    .stream()
                    .filter(page -> page.getNumber() == pageIndex)
                    .findFirst().orElseThrow(() -> new NoSuchElementException(format("ClassificationPage with number %d not found", pageIndex)));
        }

    }

}
