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

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

import java.util.ArrayList;
import java.util.List;
import java.util.ListIterator;
import java.util.Locale;

import org.apache.commons.text.similarity.LevenshteinDistance;
import org.springframework.stereotype.Service;

import com.iqser.red.service.redaction.v1.server.model.document.nodes.LayoutEngine;
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.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.Word;
import com.knecon.fforesight.service.layoutparser.processor.utils.PageInformation;
import com.knecon.fforesight.service.layoutparser.processor.utils.TextNormalizationUtilities;

import lombok.Data;

@Service
public class BlockificationPostprocessingService {

    private static final float STRING_SIMILARITY_THRESHOLD = 0.1f;


    public void findHeadlinesFromOutline(ClassificationDocument classificationDocument, int pageNumber, ClassificationPage classificationPage, PageInformation pageInformation) {

        OutlineObject lastProcessedOutlineObject = null;
        List<OutlineObject> outlineObjects = classificationDocument.getOutlineObjectTree().getOutlineObjectsPerPage().getOrDefault(pageNumber, new ArrayList<>());

        OutlineObject notFoundOutlineObject = null;
        if (lastProcessedOutlineObject != null && !lastProcessedOutlineObject.isFound()) {
            lastProcessedOutlineObject.resetPoint();
            notFoundOutlineObject = lastProcessedOutlineObject;
        }
        if (!outlineObjects.isEmpty()) {
            classificationPage.setOutlineObjects(outlineObjects);
            lastProcessedOutlineObject = sanitizeOutlineBlocks(classificationPage, notFoundOutlineObject);
        }
        classificationDocument.getLayoutDebugLayer().addOutlineObjects(outlineObjects, pageInformation);
    }


    public OutlineObject sanitizeOutlineBlocks(ClassificationPage classificationPage, OutlineObject notFoundOutlineObject) {

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

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

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

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

            OutlineObject firstOutlineObject = null;
            OutlineProcessionContext firstOutlineObjectProcessionContext = null;
            if (outlineObjectListIterator.hasNext()) {
                firstOutlineObject = outlineObjectListIterator.next();
                firstOutlineObjectProcessionContext = new OutlineProcessionContext(firstOutlineObject);
                processTextBlocks(getTextPageBlocks(classificationPage), 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), firstOutlineObjectProcessionContext);
                firstOutlineObject.setFound(selectMatch(classificationPage, firstOutlineObjectProcessionContext));
            }

        }

        outlineObjectListIterator.forEachRemaining(outlineObject -> {
            OutlineProcessionContext outlineObjectProcessionContext = new OutlineProcessionContext(outlineObject);
            processTextBlocks(getTextPageBlocks(classificationPage), 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, OutlineProcessionContext context) {

        OutlineObject outlineObject = context.getOutlineObject();
        ListIterator<TextPageBlock> iterator = textBlocks.listIterator();
        while (iterator.hasNext()) {
            TextPageBlock pageBlock = iterator.next();
            if (outlineObject.isAbove(pageBlock)) {
                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.getWords(), headline);
        if (wordSequenceResult.inSequence.isEmpty() && !headline.equals(title)) {
            wordSequenceResult = findWordSequence(blockToSplit.getWords(), title);
        }

        boolean modifiedBlockToSplit = false;
        if (!wordSequenceResult.inSequence.isEmpty()) {
            blockToSplit.setWords(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<Word> words, String text) {

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

        if (target.isBlank()) {
            return new WordSequenceResult();
        }

        for (Word sequence : words) {

            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());

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

                        Word 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(words.subList(words.indexOf(sequence) + 1, words.size()));
                    return new WordSequenceResult(inSequence, preSequence, postSequence);
                }
            }
        }

        return new WordSequenceResult();
    }


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

        Word in = null;
        Word out;

        String currentSequence = sequence.toString().toLowerCase(Locale.ROOT);
        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.length()) {
            in = createSubSequence(sequence, endIndex, sequence.length());
        }

        return new SplitSequenceResult(in, out);
    }


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

        Word newSeq = new Word(new ArrayList<>(sequence.getCharacters().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.getWords().isEmpty()) {

                    if (textPageBlock.getDir() == firstBlock.getDir()) {
                        firstBlock.addAll(textPageBlock.getWords());
                        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) {

        return outlineObject.distance(pageBlock);
    }


    // 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());

        int threshold = (int) (Math.min(blockText.length(), outlineTitle.length()) * STRING_SIMILARITY_THRESHOLD) + 1;
        int distance = new LevenshteinDistance(threshold).apply(blockText, outlineTitle);
        if (distance >= 0 && distance < threshold) {
            context.directMatch = pageBlock;
            return true;
        }

        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 TextNormalizationUtilities.removeAllWhitespaces(text)//
                .trim() // sometimes there are trailing empty bytes at the end of the string trim() seems to remove them
                .toLowerCase(Locale.ENGLISH);
    }


    @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<Word> inSequence;
        public List<Word> preSequence;
        public List<Word> postSequence;


        public WordSequenceResult(List<Word> inSequence, List<Word> preSequence, List<Word> 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(Word in, Word out) {

    }

}
