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

import static com.knecon.fforesight.service.layoutparser.processor.services.blockification.DocstrumBlockificationService.buildTextBlock;

import java.awt.geom.Rectangle2D;
import java.util.ArrayList;
import java.util.List;
import java.util.ListIterator;
import java.util.Locale;
import java.util.function.Function;

import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;

import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.LayoutEngine;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationPage;
import com.knecon.fforesight.service.layoutparser.processor.model.PageBlockType;
import com.knecon.fforesight.service.layoutparser.processor.model.SectionIdentifier;
import com.knecon.fforesight.service.layoutparser.processor.model.outline.OutlineObject;
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.services.factory.SearchTextWithTextPositionFactory;
import com.knecon.fforesight.service.layoutparser.processor.utils.RectangleTransformations;

import lombok.Data;

@Service
public class BlockificationPostprocessingService {

    private static final float BLOCK_TO_OUTLINE_DISTANCE_THRESHOLD = 5.0f;

    private static final Function<TextPageBlock, Rectangle2D> blockToBoundingBox = (abstractPageBlock) -> abstractPageBlock.getSequences()
            .stream()
            .map(textPositionSequence -> textPositionSequence.getTextPositions()
                    .stream()
                    .map(tp -> SearchTextWithTextPositionFactory.mapRedTextPositionToInitialUserSpace(tp, textPositionSequence))
                    .collect(RectangleTransformations.collectBBox()))
            .collect(RectangleTransformations.collectBBox());


    public OutlineObject sanitizeOutlineBlocks(ClassificationPage classificationPage, OutlineObject notFoundOutlineObject) {

        List<OutlineObject> outlineObjects = classificationPage.getOutlineObjects();

        if (getTextPageBlocks(classificationPage).isEmpty() || outlineObjects.isEmpty()) {
            return null;
        }

        float pageHeight = classificationPage.getPageHeight();

        ListIterator<OutlineObject> outlineObjectListIterator = outlineObjects.listIterator();

        if (notFoundOutlineObject != null) {
            OutlineProcessionContext notFoundOutlineObjectProcessionContext = new OutlineProcessionContext(notFoundOutlineObject);
            processTextBlocks(getTextPageBlocks(classificationPage), pageHeight, notFoundOutlineObjectProcessionContext);

            OutlineObject firstOutlineObject = null;
            OutlineProcessionContext firstOutlineObjectProcessionContext = null;
            if (outlineObjectListIterator.hasNext()) {
                firstOutlineObject = outlineObjectListIterator.next();
                firstOutlineObjectProcessionContext = new OutlineProcessionContext(firstOutlineObject);
                processTextBlocks(getTextPageBlocks(classificationPage), pageHeight, firstOutlineObjectProcessionContext);
            }

            if (!contextsOverlap(notFoundOutlineObjectProcessionContext, firstOutlineObjectProcessionContext)) {
                notFoundOutlineObject.setFound(selectMatch(classificationPage, notFoundOutlineObjectProcessionContext));
            }
            if (firstOutlineObject != null) {
                // re-create the context for the updated blocks
                firstOutlineObjectProcessionContext = new OutlineProcessionContext(firstOutlineObject);
                processTextBlocks(getTextPageBlocks(classificationPage), pageHeight, firstOutlineObjectProcessionContext);
                firstOutlineObject.setFound(selectMatch(classificationPage, firstOutlineObjectProcessionContext));
            }

        }

        outlineObjectListIterator.forEachRemaining(outlineObject -> {
            OutlineProcessionContext outlineObjectProcessionContext = new OutlineProcessionContext(outlineObject);
            processTextBlocks(getTextPageBlocks(classificationPage), pageHeight, outlineObjectProcessionContext);
            outlineObject.setFound(selectMatch(classificationPage, outlineObjectProcessionContext));
        });

        if (!outlineObjects.isEmpty()) {
            return outlineObjects.get(outlineObjects.size() - 1);
        } else {
            return notFoundOutlineObject;
        }
    }


    private static List<TextPageBlock> getTextPageBlocks(ClassificationPage classificationPage) {

        return classificationPage.getTextBlocks()
                .stream()
                .filter(block -> block instanceof TextPageBlock)
                .map(block -> (TextPageBlock) block)
                .toList();
    }


    private boolean contextsOverlap(OutlineProcessionContext notFoundOutlineObjectProcessionContext, OutlineProcessionContext firstOutlineObjectProcessionContext) {

        if (firstOutlineObjectProcessionContext == null) {
            return false;
        }

        String notFoundTitle = notFoundOutlineObjectProcessionContext.getOutlineObject().getTitle();
        String firstTitle = firstOutlineObjectProcessionContext.getOutlineObject().getTitle();

        if (!firstTitle.startsWith(notFoundTitle)) {
            return false;
        }

        var blocksOfNotFoundOutline = getAllMatchingBlocks(notFoundOutlineObjectProcessionContext);
        var blocksOfFirstOutline = getAllMatchingBlocks(firstOutlineObjectProcessionContext);

        double maxYFirst = blocksOfFirstOutline.stream()
                .mapToDouble(TextPageBlock::getPdfMaxY)
                .max()
                .orElse(Double.NEGATIVE_INFINITY);

        return blocksOfNotFoundOutline.stream()
                .mapToDouble(TextPageBlock::getPdfMaxY)
                .anyMatch(y -> y >= maxYFirst);
    }


    private List<TextPageBlock> getAllMatchingBlocks(OutlineProcessionContext context) {

        List<TextPageBlock> blocks = new ArrayList<>();
        if (context.getDirectMatch() != null) {
            blocks.add(context.getDirectMatch());
        }
        if (context.getSplitCandidate() != null) {
            blocks.add(context.getSplitCandidate());
        }
        blocks.addAll(context.getMergeCandidates());
        return blocks;
    }


    private void processTextBlocks(List<TextPageBlock> textBlocks, float pageHeight, OutlineProcessionContext context) {

        OutlineObject outlineObject = context.getOutlineObject();
        ListIterator<TextPageBlock> iterator = textBlocks.listIterator();
        while (iterator.hasNext()) {
            TextPageBlock pageBlock = iterator.next();
            if (pageHeight - outlineObject.getPoint().getY() - BLOCK_TO_OUTLINE_DISTANCE_THRESHOLD <= pageBlock.getMaxY()) {
                break;
            }
        }
        if (iterator.hasPrevious()) {
            iterator.previous();
        }
        boolean earlyStop = false;
        while (iterator.hasNext() && !earlyStop) {
            TextPageBlock pageBlock = iterator.next();
            earlyStop = processOutlineForTextBlock(pageBlock, context);
        }
    }


    private boolean selectMatch(ClassificationPage classificationPage, OutlineProcessionContext context) {

        OutlineObject outlineObject = context.outlineObject;
        TextPageBlock directMatch = context.directMatch;
        List<TextPageBlock> mergeCandidates = context.mergeCandidates;
        TextPageBlock splitCandidate = context.splitCandidate;
        PageBlockType headlineType = PageBlockType.getHeadlineType(outlineObject.getTreeDepth());

        double distanceToDirectMatch = directMatch != null ? calculateDistance(outlineObject, directMatch) : Double.MAX_VALUE;
        double distanceToSplitCandidate = splitCandidate != null ? calculateDistance(outlineObject, splitCandidate) : Double.MAX_VALUE;

        double distanceToBestMergeCandidates = Double.MAX_VALUE;
        List<TextPageBlock> bestMergeCandidateCombination = new ArrayList<>();
        if (!mergeCandidates.isEmpty()) {

            // with this code adjacent blocks to the first and last merge candidate get added, this could be useful for some edge cases:
            //List<TextPageBlock> allMergeCandidates = new ArrayList<>(mergeCandidates);
            //addNeighborsOfCandidate(kdTree, mergeCandidates.get(0), allMergeCandidates);
            //if (mergeCandidates.size() > 1) {
            //    addNeighborsOfCandidate(kdTree, mergeCandidates.get(mergeCandidates.size() - 1), allMergeCandidates);
            //}
            //allMergeCandidates = allMergeCandidates.stream()
            //        .distinct()
            //        .toList();

            List<List<TextPageBlock>> combinations = findCombinations(outlineObject.getTitle(), mergeCandidates);

            for (List<TextPageBlock> combination : combinations) {
                double averageDistance = combination.stream()
                        .map(block -> calculateDistance(outlineObject, block))
                        .mapToDouble(Double::doubleValue).average()
                        .orElse(Double.MAX_VALUE);
                if (distanceToBestMergeCandidates > averageDistance) {
                    distanceToBestMergeCandidates = averageDistance;
                    bestMergeCandidateCombination = combination;
                }
            }
        }

        double minDistance = Math.min(distanceToDirectMatch, Math.min(distanceToSplitCandidate, distanceToBestMergeCandidates));

        if (minDistance == Double.MAX_VALUE) {
            return false;
        }
        if (minDistance == distanceToDirectMatch) {
            setClassificationAndAddOutlineEngine(directMatch, headlineType);
        } else if (minDistance == distanceToSplitCandidate) {
            SplitBlockResult splitBlockResult = splitBlock(classificationPage, splitCandidate, context.sectionIdentifier, outlineObject.getTitle());
            if (splitBlockResult.modifiedBlockToSplit) {
                setClassificationAndAddOutlineEngine(splitCandidate, headlineType);
            }
            splitBlockResult.otherBlocks.forEach(other -> other.setClassification(null));
        } else {
            var merged = mergeBlocks(classificationPage, bestMergeCandidateCombination);
            setClassificationAndAddOutlineEngine(merged, headlineType);
        }
        return true;
    }


    private static void setClassificationAndAddOutlineEngine(TextPageBlock block, PageBlockType headlineType) {

        block.setClassification(headlineType);
        block.getEngines().add(LayoutEngine.OUTLINE);
    }


    private SplitBlockResult splitBlock(ClassificationPage classificationPage, TextPageBlock blockToSplit, SectionIdentifier sectionIdentifier, String title) {

        List<TextPageBlock> otherBlocks = new ArrayList<>();
        int blockToSplitIdx = classificationPage.getTextBlocks().indexOf(blockToSplit);

        String headline = title;
        if (!sectionIdentifier.getFormat().equals(SectionIdentifier.Format.EMPTY) && !title.startsWith(sectionIdentifier.getIdentifierString())) {
            headline = sectionIdentifier + headline;
        }

        WordSequenceResult wordSequenceResult = findWordSequence(blockToSplit.getSequences(), headline);
        if (wordSequenceResult.inSequence.isEmpty() && !headline.equals(title)) {
            wordSequenceResult = findWordSequence(blockToSplit.getSequences(), title);
        }

        boolean modifiedBlockToSplit = false;
        if (!wordSequenceResult.inSequence.isEmpty()) {
            blockToSplit.setSequences(wordSequenceResult.inSequence);
            blockToSplit.recalculateBBox();
            modifiedBlockToSplit = true;
        }

        if (!wordSequenceResult.preSequence.isEmpty()) {
            TextPageBlock block = buildTextBlock(wordSequenceResult.preSequence, 0);
            classificationPage.getTextBlocks().add(blockToSplitIdx, block);
            otherBlocks.add(block);
            blockToSplitIdx++;
        }
        if (!wordSequenceResult.postSequence.isEmpty()) {
            TextPageBlock block = buildTextBlock(wordSequenceResult.postSequence, 0);
            classificationPage.getTextBlocks().add(blockToSplitIdx + 1, block);
            otherBlocks.add(block);
        }

        return new SplitBlockResult(modifiedBlockToSplit, otherBlocks);
    }


    private static WordSequenceResult findWordSequence(List<TextPositionSequence> textPositionSequences, String text) {

        String target = sanitizeString(text);
        List<TextPositionSequence> inSequence = new ArrayList<>();
        List<TextPositionSequence> preSequence = new ArrayList<>();
        List<TextPositionSequence> postSequence = new ArrayList<>();
        StringBuilder currentSequence = new StringBuilder();

        for (TextPositionSequence sequence : textPositionSequences) {

            currentSequence.append(sanitizeString(sequence.toString()));
            inSequence.add(sequence);

            if (currentSequence.length() >= target.length()) {

                if (currentSequence.toString().endsWith(target)) {

                    int index = 0;
                    String toRemove = currentSequence.substring(0, currentSequence.length() - target.length());

                    TextPositionSequence next = inSequence.get(index);
                    while (currentSequence.length() - next.length() >= target.length()) {

                        TextPositionSequence removed = inSequence.remove(index);
                        currentSequence.delete(0, removed.toString().length());
                        preSequence.add(removed);

                        next = inSequence.get(index);
                        toRemove = toRemove.substring(removed.length());
                    }

                    if (!toRemove.isEmpty()) {
                        SplitSequenceResult splitSequenceResult = splitSequence(inSequence.remove(index), toRemove);

                        currentSequence.delete(0, splitSequenceResult.out.length());
                        preSequence.add(splitSequenceResult.out);
                        inSequence.add(index, splitSequenceResult.in);
                    }

                } else if (currentSequence.toString().startsWith(target)) {

                    int index = inSequence.size() - 1;
                    String toRemove = currentSequence.substring(target.length());

                    SplitSequenceResult splitSequenceResult = splitSequence(inSequence.remove(index), toRemove);
                    currentSequence.delete(currentSequence.length() - splitSequenceResult.out.length(), currentSequence.length());

                    inSequence.add(index, splitSequenceResult.in);
                    postSequence.add(splitSequenceResult.out);
                }

                if (currentSequence.toString().equals(target)) {
                    postSequence.addAll(textPositionSequences.subList(textPositionSequences.indexOf(sequence) + 1, textPositionSequences.size()));
                    return new WordSequenceResult(inSequence, preSequence, postSequence);
                }
            }
        }

        return new WordSequenceResult();
    }


    private static SplitSequenceResult splitSequence(TextPositionSequence sequence, String toRemove) {

        TextPositionSequence in = null;
        TextPositionSequence out;

        String currentSequence = sequence.toString();
        int index = currentSequence.indexOf(toRemove);
        int endIndex = index + toRemove.length();

        out = createSubSequence(sequence, index, endIndex);

        if (index > 0) {
            in = createSubSequence(sequence, 0, index);
        } else if (endIndex < sequence.getTextPositions().size()) {
            in = createSubSequence(sequence, endIndex, sequence.getTextPositions().size());
        }

        return new SplitSequenceResult(in, out);
    }


    private static TextPositionSequence createSubSequence(TextPositionSequence sequence, int start, int end) {

        TextPositionSequence newSeq = new TextPositionSequence(new ArrayList<>(sequence.getTextPositions().subList(start, end)), sequence.getPage());
        newSeq.setParagraphStart(sequence.isParagraphStart());
        return newSeq;
    }


    private TextPageBlock mergeBlocks(ClassificationPage classificationPage, List<TextPageBlock> blocksToMerge) {

        TextPageBlock firstBlock = blocksToMerge.get(0);

        if (blocksToMerge.size() > 1) {

            List<TextPageBlock> mergedBlocks = new ArrayList<>();
            for (TextPageBlock textPageBlock : blocksToMerge.subList(1, blocksToMerge.size())) {

                if (firstBlock != null && !firstBlock.getSequences().isEmpty()) {

                    if (textPageBlock.getDir() == firstBlock.getDir()) {
                        firstBlock.getSequences().addAll(textPageBlock.getSequences());
                        mergedBlocks.add(textPageBlock);
                    }
                }
            }

            assert firstBlock != null;
            firstBlock.setToDuplicate(false);
            firstBlock.recalculateBBox();
            classificationPage.getTextBlocks().removeAll(mergedBlocks);
        }

        return firstBlock;
    }


    private static List<List<TextPageBlock>> findCombinations(String title, List<TextPageBlock> blocks) {

        List<List<TextPageBlock>> combinations = new ArrayList<>();
        findCombinations(title, blocks, new ArrayList<>(), combinations);
        return combinations;
    }


    private static void findCombinations(String title, List<TextPageBlock> blocks, List<TextPageBlock> current, List<List<TextPageBlock>> combinations) {

        String target = sanitizeString(title);
        if (target.isEmpty()) {
            combinations.add(new ArrayList<>(current));
            return;
        }

        List<TextPageBlock> remaining = blocks.stream()
                .filter(block -> !current.contains(block))
                .toList();
        for (TextPageBlock block : remaining) {
            String prefix = sanitizeString(block.getText());
            if (target.startsWith(prefix)) {
                current.add(block);
                findCombinations(target.substring(prefix.length()), blocks.subList(blocks.indexOf(block) + 1, blocks.size()), current, combinations);
                current.remove(current.size() - 1);
            }
        }
    }


    private double calculateDistance(OutlineObject outlineObject, TextPageBlock pageBlock) {

        double deltaX = outlineObject.getPoint().getX() - pageBlock.getMinX();
        double deltaY = pageBlock.getPageHeight() - outlineObject.getPoint().getY() - pageBlock.getMinY();
        return Math.sqrt(deltaX * deltaX + deltaY * deltaY);
    }


    // currently only three cases are handled here:
    // 1. equality
    // 2. outline title contains block text
    // 3. block text contains outline title
    // another possible case is an intersection, meaning a title is split up between two different blocks
    // this should not happen with how docstrum creates the blocks
    // if it is indeed necessary, a splitting has to be done with a follow-up merge
    private boolean processOutlineForTextBlock(TextPageBlock pageBlock, OutlineProcessionContext context) {

        OutlineObject outlineObject = context.getOutlineObject();
        String blockText = sanitizeString(pageBlock.getText());
        String outlineTitle = sanitizeString(outlineObject.getTitle());

        boolean blockTextContainsOutlineTitle = blockText.contains(outlineTitle);
        boolean outlineTitleContainsBlockText = outlineTitle.contains(blockText);

        if (!blockTextContainsOutlineTitle && !outlineTitleContainsBlockText) {
            return false;
        }

        if (blockText.equals(outlineTitle) && context.directMatch == null) {
            context.directMatch = pageBlock;
            return true;
        }

        if (outlineTitleContainsBlockText) {
            context.mergeCandidates.add(pageBlock);
        }

        if (blockTextContainsOutlineTitle) {
            SectionIdentifier sectionIdentifier = SectionIdentifier.fromSearchText(blockText);

            if (sectionIdentifier.getFormat() != SectionIdentifier.Format.EMPTY && !outlineTitle.startsWith(sectionIdentifier.getIdentifierString())) {

                if (blockText.startsWith(sectionIdentifier.getIdentifierString()) && blockText.endsWith(outlineTitle) && context.directMatch == null) {
                    context.directMatch = pageBlock;
                    return true;
                } else if (context.splitCandidate == null) {
                    context.sectionIdentifier = sectionIdentifier;
                }
            }
            if (context.splitCandidate == null) {
                context.splitCandidate = pageBlock;
            }
        }
        return false;
    }


    private static String sanitizeString(String text) {

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


    @Data
    private static class OutlineProcessionContext {

        private TextPageBlock directMatch;
        private OutlineObject outlineObject;
        private List<TextPageBlock> mergeCandidates;
        private TextPageBlock splitCandidate;
        private SectionIdentifier sectionIdentifier;


        OutlineProcessionContext(OutlineObject outlineObject) {

            this.outlineObject = outlineObject;
            this.directMatch = null;
            this.mergeCandidates = new ArrayList<>();
            this.splitCandidate = null;
            this.sectionIdentifier = SectionIdentifier.empty();
        }

    }

    public static class WordSequenceResult {

        public List<TextPositionSequence> inSequence;
        public List<TextPositionSequence> preSequence;
        public List<TextPositionSequence> postSequence;


        public WordSequenceResult(List<TextPositionSequence> inSequence, List<TextPositionSequence> preSequence, List<TextPositionSequence> postSequence) {

            this.inSequence = inSequence;
            this.preSequence = preSequence;
            this.postSequence = postSequence;
        }


        public WordSequenceResult() {

            this.inSequence = new ArrayList<>();
            this.preSequence = new ArrayList<>();
            this.postSequence = new ArrayList<>();
        }

    }

    public record SplitBlockResult(boolean modifiedBlockToSplit, List<TextPageBlock> otherBlocks) {

    }

    public record SplitSequenceResult(TextPositionSequence in, TextPositionSequence out) {

    }

}
