package com.knecon.fforesight.service.layoutparser.processor.model.outline;

import static com.knecon.fforesight.service.layoutparser.processor.model.PageBlockType.TABLE_OF_CONTENTS_HEADLINE;
import static com.knecon.fforesight.service.layoutparser.processor.model.PageBlockType.getHeadlineNumber;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeSet;

import org.springframework.stereotype.Service;

import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationDocument;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPageBlock;

import io.micrometer.observation.annotation.Observed;
import lombok.extern.slf4j.Slf4j;

@Service
@Slf4j
public class SectionTreeBuilderService {

    @Observed(name = "OutlineValidationService", contextualName = "create-toc")
    public SectionTree createSectionTree(ClassificationDocument classificationDocument) {

        List<TextPageBlock> headlines = extractHeadlines(classificationDocument);

        List<SectionTreeEntry> mainSections = new ArrayList<>();
        Map<Integer, SectionTreeEntry> lastItemsPerDepth = new HashMap<>();
        SectionTreeEntry last = null;
        TreeSet<Integer> depths = new TreeSet<>();

        for (TextPageBlock current : headlines) {
            int currentDepth = getHeadlineNumber(current.getClassification());
            Integer parentDepth = depths.floor(currentDepth - 1);

            var tocItem = new SectionTreeEntry(current);

            if (parentDepth == null) {
                mainSections.add(tocItem);
                lastItemsPerDepth = new HashMap<>();
                depths = new TreeSet<>();

            } else {
                assert last != null;
                int lastDepth = getHeadlineNumber(last.getHeadline().getClassification());
                if (last.getHeadline().getClassification().equals(TABLE_OF_CONTENTS_HEADLINE) && !current.getClassification().equals(TABLE_OF_CONTENTS_HEADLINE)) {
                    // headline after toc should always start a main section
                    parentDepth = 1;
                } else if (lastDepth < parentDepth) {
                    parentDepth = lastDepth;
                } else if (lastDepth == currentDepth && last.getParent() != null) {
                    parentDepth = getHeadlineNumber(last.getParent().getHeadline().getClassification());
                }

                SectionTreeEntry parent = lastItemsPerDepth.get(parentDepth);
                parent.addChild(tocItem);
            }

            last = tocItem;
            lastItemsPerDepth.put(currentDepth, tocItem);
            depths.add(currentDepth);
        }

        return new SectionTree(mainSections);

    }


    private static List<TextPageBlock> extractHeadlines(ClassificationDocument classificationDocument) {

        return classificationDocument.getPages()
                .stream()
                .flatMap(classificationPage -> classificationPage.getTextBlocks()
                        .stream()
                        .filter(tb -> tb instanceof TextPageBlock && tb.getClassification() != null && tb.getClassification().isHeadline())
                        .map(tb -> (TextPageBlock) tb))
                .toList();
    }

}
