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

import java.util.Arrays;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;

import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentData;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentPage;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentPositionData;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentStructure;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentTextData;
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.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.nodes.SemanticNode;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Table;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.TableCell;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.textblock.AtomicTextBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.textblock.TextBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.textblock.TextBlockCollector;

import lombok.experimental.UtilityClass;

@UtilityClass
public class DocumentGraphMapper {

    public Document toDocumentGraph(DocumentData documentData) {

        Document document = new Document();
        DocumentTree documentTree = new DocumentTree(document);
        Context context = new Context(documentData, documentTree);

        context.pages.addAll(Arrays.stream(documentData.getDocumentPages()).map(DocumentGraphMapper::buildPage).toList());

        context.documentTree.getRoot().getChildren().addAll(buildEntries(documentData.getDocumentStructure().getRoot().getChildren(), context));

        document.setDocumentTree(context.documentTree);
        document.setPages(new HashSet<>(context.pages));
        document.setNumberOfPages(documentData.getDocumentPages().length);

        document.setTextBlock(document.getTextBlock());
        return document;
    }


    private List<DocumentTree.Entry> buildEntries(List<DocumentStructure.EntryData> entries, Context context) {

        List<DocumentTree.Entry> newEntries = new LinkedList<>();
        for (DocumentStructure.EntryData entryData : entries) {

            List<Page> pages = Arrays.stream(entryData.getPageNumbers()).map(pageNumber -> getPage(pageNumber, context)).toList();

            SemanticNode node = switch (entryData.getType()) {
                case SECTION -> buildSection(context);
                case PARAGRAPH -> buildParagraph(context, entryData.getProperties());
                case HEADLINE -> buildHeadline(context);
                case HEADER -> buildHeader(context);
                case FOOTER -> buildFooter(context);
                case TABLE -> buildTable(context, entryData.getProperties());
                case TABLE_CELL -> buildTableCell(context, entryData.getProperties());
                case IMAGE -> buildImage(context, entryData.getProperties(), entryData.getPageNumbers());
                default -> throw new UnsupportedOperationException("Not yet implemented for type " + entryData.getType());
            };

            if (entryData.getAtomicBlockIds().length > 0) {
                TextBlock textBlock = toTextBlock(entryData.getAtomicBlockIds(), context, node);
                node.setLeafTextBlock(textBlock);
            }
            List<Integer> treeId = Arrays.stream(entryData.getTreeId()).boxed().toList();
            node.setTreeId(treeId);

            switch (entryData.getType()) {
                case HEADER -> pages.forEach(page -> page.setHeader((Header) node));
                case FOOTER -> pages.forEach(page -> page.setFooter((Footer) node));
                default -> pages.forEach(page -> page.getMainBody().add(node));
            }

            newEntries.add(DocumentTree.Entry.builder().treeId(treeId).children(buildEntries(entryData.getChildren(), context)).node(node).build());
        }
        return newEntries;
    }


    private Headline buildHeadline(Context context) {

        return Headline.builder().documentTree(context.documentTree).build();
    }


    private Image buildImage(Context context, Map<String, String> properties, Long[] pageNumbers) {

        assert pageNumbers.length == 1;
        Page page = getPage(pageNumbers[0], context);
        var builder = Image.builder();
        PropertiesMapper.parseImageProperties(properties, builder);
        return builder.documentTree(context.documentTree).page(page).build();
    }


    private TableCell buildTableCell(Context context, Map<String, String> properties) {

        TableCell.TableCellBuilder builder = TableCell.builder();
        PropertiesMapper.parseTableCellProperties(properties, builder);
        return builder.documentTree(context.documentTree).build();
    }


    private Table buildTable(Context context, Map<String, String> properties) {

        Table.TableBuilder builder = Table.builder();
        PropertiesMapper.parseTableProperties(properties, builder);
        return builder.documentTree(context.documentTree).build();
    }


    private Footer buildFooter(Context context) {

        return Footer.builder().documentTree(context.documentTree).build();
    }


    private Header buildHeader(Context context) {

        return Header.builder().documentTree(context.documentTree).build();
    }


    private Section buildSection(Context context) {

        return Section.builder().documentTree(context.documentTree).build();
    }


    private Paragraph buildParagraph(Context context, Map<String, String> properties) {

        if (PropertiesMapper.isDuplicateParagraph(properties)) {

            DuplicatedParagraph duplicatedParagraph = DuplicatedParagraph.builder().documentTree(context.documentTree).build();

            Long[] unsortedTextblockIds = PropertiesMapper.getUnsortedTextblockIds(properties);
            duplicatedParagraph.setUnsortedLeafTextBlock(toTextBlock(unsortedTextblockIds, context, duplicatedParagraph));
            return duplicatedParagraph;

        }

        return Paragraph.builder().documentTree(context.documentTree).build();
    }


    private TextBlock toTextBlock(Long[] atomicTextBlockIds, Context context, SemanticNode parent) {

        return Arrays.stream(atomicTextBlockIds).map(atomicTextBlockId -> getAtomicTextBlock(context, parent, atomicTextBlockId)).collect(new TextBlockCollector());
    }


    private AtomicTextBlock getAtomicTextBlock(Context context, SemanticNode parent, Long atomicTextBlockId) {

        return AtomicTextBlock.fromAtomicTextBlockData(context.documentTextDataBlockData.get(Math.toIntExact(atomicTextBlockId)),
                context.atomicPositionBlockData.get(Math.toIntExact(atomicTextBlockId)),
                parent,
                getPage(context.documentTextDataBlockData.get(Math.toIntExact(atomicTextBlockId)).getPage(), context));
    }


    private Page buildPage(DocumentPage p) {

        return Page.builder().rotation(p.getRotation()).height(p.getHeight()).width(p.getWidth()).number(p.getNumber()).mainBody(new LinkedList<>()).build();
    }


    private Page getPage(Long pageIndex, Context context) {

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


    static final class Context {

        private final DocumentTree documentTree;
        private final List<Page> pages;
        private final List<DocumentTextData> documentTextDataBlockData;
        private final List<DocumentPositionData> atomicPositionBlockData;


        Context(DocumentData documentData, DocumentTree documentTree) {

            this.documentTree = documentTree;
            this.pages = new LinkedList<>();
            this.documentTextDataBlockData = Arrays.stream(documentData.getDocumentTextData()).toList();
            this.atomicPositionBlockData = Arrays.stream(documentData.getDocumentPositions()).toList();

        }

    }

}
