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

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

import java.util.ArrayList;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

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.graph.nodes.Document;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.GenericSemanticNode;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Page;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Section;
import com.knecon.fforesight.service.layoutparser.processor.model.image.ClassifiedImage;
import com.knecon.fforesight.service.layoutparser.processor.model.table.TablePageBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.utils.TableMergingUtility;

import lombok.experimental.UtilityClass;

@UtilityClass
public class SectionNodeFactory {

    public void addSection(LayoutParsingType layoutParsingType,
                           GenericSemanticNode parentNode,
                           List<AbstractPageBlock> pageBlocks,
                           List<ClassifiedImage> images,
                           DocumentGraphFactory.Context context,
                           Document document) {

        // This is for the case where we have images on a page without any text/footer/header.
        // The pageBlocks list is empty, but we still need to add those images to the document.
        if (!images.isEmpty() && pageBlocks.isEmpty()) {
            images.stream()
                    .distinct()
                    .forEach(image -> DocumentGraphFactory.addImage(document, image, context));
            return;
        }

        if (pageBlocks.isEmpty()) {
            return;
        }

        Map<Integer, List<AbstractPageBlock>> blocksPerPage = pageBlocks.stream()
                .collect(groupingBy(AbstractPageBlock::getPage));
        Section section = Section.builder().documentTree(context.getDocumentTree()).build();

        context.getSections().add(section);
        blocksPerPage.keySet()
                .forEach(pageNumber -> addSectionNodeToPageNode(context, section, pageNumber));

        section.setTreeId(getTreeId(parentNode, context, section));

        addFirstHeadlineDirectlyToSection(layoutParsingType, pageBlocks, context, section, document);
        if (containsTablesAndTextBlocks(pageBlocks)) {
            splitPageBlocksIntoSubSections(pageBlocks).forEach(subSectionPageBlocks -> addSection(layoutParsingType,
                                                                                                  section,
                                                                                                  subSectionPageBlocks,
                                                                                                  emptyList(),
                                                                                                  context,
                                                                                                  document));
        } else {
            addTablesAndParagraphsAndHeadlinesToSection(layoutParsingType, pageBlocks, context, section, document);
        }

        images.stream()
                .distinct()
                .forEach(image -> DocumentGraphFactory.addImage(section, image, context));
    }


    private List<Integer> getTreeId(GenericSemanticNode parentNode, DocumentGraphFactory.Context context, Section section) {

        if (parentNode == null) {
            return context.getDocumentTree().createNewMainEntryAndReturnId(section);
        } else {
            return context.getDocumentTree().createNewChildEntryAndReturnId(parentNode, section);
        }
    }


    private void addFirstHeadlineDirectlyToSection(LayoutParsingType layoutParsingType,
                                                   List<AbstractPageBlock> pageBlocks,
                                                   DocumentGraphFactory.Context context,
                                                   Section section,
                                                   Document document) {

        if (pageBlocks.get(0).isHeadline()) {
            addTablesAndParagraphsAndHeadlinesToSection(layoutParsingType, List.of(pageBlocks.get(0)), context, section, document);
            pageBlocks.remove(0);
        }
    }


    private void addTablesAndParagraphsAndHeadlinesToSection(LayoutParsingType layoutParsingType,
                                                             List<AbstractPageBlock> pageBlocks,
                                                             DocumentGraphFactory.Context context,
                                                             Section section,
                                                             Document document) {

        Set<AbstractPageBlock> alreadyMerged = new HashSet<>();
        List<AbstractPageBlock> remainingBlocks = new LinkedList<>(pageBlocks);
        for (AbstractPageBlock abstractPageBlock : pageBlocks) {

            if (alreadyMerged.contains(abstractPageBlock)) {
                continue;
            }

            remainingBlocks.removeAll(alreadyMerged);

            if (abstractPageBlock instanceof TextPageBlock) {

                switch (layoutParsingType) {
                    case REDACT_MANAGER, CLARIFYND, CLARIFYND_PARAGRAPH_DEBUG -> {
                        alreadyMerged.add(abstractPageBlock);
                        remainingBlocks.remove(abstractPageBlock);
                        DocumentGraphFactory.addParagraphOrHeadline(section, (TextPageBlock) abstractPageBlock, context, new ArrayList<>(), layoutParsingType);
                    }
                    default -> {
                        List<TextPageBlock> textBlocks = findTextBlocksWithSameClassificationAndAlignsY((TextPageBlock) abstractPageBlock, remainingBlocks);
                        alreadyMerged.addAll(textBlocks);
                        DocumentGraphFactory.addParagraphOrHeadline(section, (TextPageBlock) abstractPageBlock, context, textBlocks, layoutParsingType);
                    }
                }
            } else if (abstractPageBlock instanceof TablePageBlock tablePageBlock) {
                List<TablePageBlock> tablesToMerge = TableMergingUtility.findConsecutiveTablesWithSameColCountAndSameHeaders(tablePageBlock, remainingBlocks);
                alreadyMerged.addAll(tablesToMerge);
                TableNodeFactory.addTable(layoutParsingType, section, tablesToMerge, context, document);
            } else {
                throw new RuntimeException(format("Unhandled AbstractPageBlockType %s!", abstractPageBlock.getClass()));
            }
        }
    }


    private boolean containsTablesAndTextBlocks(List<AbstractPageBlock> pageBlocks) {

        return pageBlocks.stream()
                       .anyMatch(pageBlock -> pageBlock instanceof TablePageBlock) && pageBlocks.stream()
                       .anyMatch(pageBlock -> pageBlock instanceof TextPageBlock);
    }


    /**
     * This function splits the list of PageBlocks around TablePageBlocks, such that SubSections can be created, that don't include tables.
     * This is needed so we can execute rules on sections, that do not contain tables.
     * See: <a href="https://knecon.atlassian.net/wiki/spaces/RED/pages/14765218/Document+Structure">document structure wiki</a>
     *
     * @param pageBlocks a List of AbstractPageBlocks, which have at least one TablePageBlock and one ClassificationTextBlock
     * @return List of Lists of AbstractPageBlocks, which include either a single Headline ClassificationTextBlock and a TablePageBlock or only ClassificationTextBlocks.
     */
    private List<List<AbstractPageBlock>> splitPageBlocksIntoSubSections(List<AbstractPageBlock> pageBlocks) {

        List<List<AbstractPageBlock>> splitList = splitIntoCoherentList(pageBlocks);
        movePrecedingHeadlineToTableList(splitList);
        return splitList.stream()
                .filter(list -> !list.isEmpty())
                .toList();
    }


    private void movePrecedingHeadlineToTableList(List<List<AbstractPageBlock>> splitList) {

        for (int i = 0; i < splitList.size(); i++) {
            if (listIsTablesOnly(splitList.get(i)) && i > 0) {
                List<AbstractPageBlock> previousList = splitList.get(i - 1);
                AbstractPageBlock lastPageBlockInPreviousList = previousList.get(previousList.size() - 1);
                if (lastPageBlockInPreviousList.isHeadline()) {
                    previousList.remove(previousList.size() - 1);
                    splitList.get(i).add(0, lastPageBlockInPreviousList);
                }
            }
        }
    }


    private boolean listIsTablesOnly(List<AbstractPageBlock> abstractPageBlocks) {

        return abstractPageBlocks.stream()
                .allMatch(abstractPageBlock -> abstractPageBlock instanceof TablePageBlock);
    }


    /**
     * @param pageBlocks a List of AbstractPageBlocks, which have at least one TablePageBlock and one ClassificationTextBlock
     * @return List of Lists of AbstractPageBlocks, which are exclusively of type ClassificationTextBlock or TablePageBlock
     */
    private List<List<AbstractPageBlock>> splitIntoCoherentList(List<AbstractPageBlock> pageBlocks) {

        List<List<AbstractPageBlock>> splitList = new LinkedList<>();
        List<AbstractPageBlock> currentList = new LinkedList<>();
        splitList.add(currentList);

        Class<? extends AbstractPageBlock> lastPageBlockClass = pageBlocks.get(0).getClass();
        for (AbstractPageBlock pageBlock : pageBlocks) {
            if (lastPageBlockClass.isInstance(pageBlock)) {
                currentList.add(pageBlock);
            } else {
                currentList = new LinkedList<>();
                currentList.add(pageBlock);
                splitList.add(currentList);
                lastPageBlockClass = pageBlock.getClass();
            }
        }
        return splitList;
    }


    private List<TextPageBlock> findTextBlocksWithSameClassificationAndAlignsY(TextPageBlock atc, List<AbstractPageBlock> pageBlocks) {

        return pageBlocks.stream()
                .filter(abstractTextContainer -> !abstractTextContainer.equals(atc))
                .filter(abstractTextContainer -> abstractTextContainer.getPage() == atc.getPage())
                .filter(abstractTextContainer -> abstractTextContainer instanceof TextPageBlock)
                .filter(abstractTextContainer -> abstractTextContainer.intersectsY(atc))
                .map(abstractTextContainer -> (TextPageBlock) abstractTextContainer)
                .filter(abstractTextContainer -> abstractTextContainer.getDir() == atc.getDir())
                .filter(abstractTextContainer -> !abstractTextContainer.isToDuplicate())
                .toList();
    }


    private void addSectionNodeToPageNode(DocumentGraphFactory.Context context, Section section, Integer pageNumber) {

        Page page = context.getPage(pageNumber);
        page.getMainBody().add(section);
    }

}
