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

import java.awt.geom.Rectangle2D;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

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.internal.api.data.redaction.LayoutEngine;
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.Image;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Page;
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 lombok.experimental.UtilityClass;

@UtilityClass
public class DocumentDataMapper {

    public DocumentData toDocumentData(Document document) {

        List<DocumentTextData> documentTextData = document.streamTerminalTextBlocksInOrder()
                .flatMap(textBlock -> textBlock.getAtomicTextBlocks().stream())
                .distinct()
                .map(DocumentDataMapper::toAtomicTextBlockData)
                .toList();

        List<DocumentPositionData> atomicPositionBlockData = document.streamTerminalTextBlocksInOrder()
                .flatMap(textBlock -> textBlock.getAtomicTextBlocks().stream())
                .distinct()
                .map(DocumentDataMapper::toAtomicPositionBlockData)
                .toList();

        Set<Long> nonEmptyTextBlocks = documentTextData.stream().mapToLong(DocumentTextData::getId).boxed().collect(Collectors.toSet());

        List<DocumentPage> documentPageData = document.getPages().stream().map(DocumentDataMapper::toPageData).toList();
        DocumentStructure tableOfContentsData = toDocumentTreeData(document.getDocumentTree());
        return DocumentData.builder()
                .documentTextData(documentTextData.toArray(new DocumentTextData[0]))
                .documentPositions(atomicPositionBlockData.toArray(new DocumentPositionData[0]))
                .documentPages(documentPageData.toArray(new DocumentPage[0]))
                .documentStructure(tableOfContentsData)
                .build();
    }


    private DocumentStructure toDocumentTreeData(DocumentTree documentTree) {

        return new DocumentStructure(toEntryData(documentTree.getRoot()));
    }


    private DocumentStructure.EntryData toEntryData(DocumentTree.Entry entry) {

        Long[] atomicTextBlocks;

        if (entry.getNode().isLeaf()) {
            atomicTextBlocks = toAtomicTextBlockIds(entry.getNode().getLeafTextBlock());
        } else {
            atomicTextBlocks = new Long[]{};
        }

        Map<String, String> properties = switch (entry.getType()) {
            case TABLE -> PropertiesMapper.buildTableProperties((Table) entry.getNode());
            case TABLE_CELL -> PropertiesMapper.buildTableCellProperties((TableCell) entry.getNode());
            case IMAGE -> PropertiesMapper.buildImageProperties((Image) entry.getNode());
            case PARAGRAPH ->
                    entry.getNode() instanceof DuplicatedParagraph duplicatedParagraph ? PropertiesMapper.buildDuplicateParagraphProperties(duplicatedParagraph) : new HashMap<>();
            default -> new HashMap<>();
        };

        DocumentStructure.EntryData.EntryDataBuilder documentBuilder = DocumentStructure.EntryData.builder()
                .treeId(toPrimitiveIntArray(entry.getTreeId()))
                .children(entry.getChildren().stream().map(DocumentDataMapper::toEntryData).toList())
                .type(entry.getType())
                .atomicBlockIds(atomicTextBlocks)
                .pageNumbers(entry.getNode().getPages().stream().map(Page::getNumber).map(Integer::longValue).toArray(Long[]::new))
                .properties(properties);
        if (entry.getNode() != null) {
            documentBuilder.engines(entry.getNode().getEngines());
        } else {
            documentBuilder.engines(new HashSet<>(Set.of(LayoutEngine.ALGORITHM)));
        }
        return documentBuilder.build();
    }


    private Long[] toAtomicTextBlockIds(TextBlock textBlock) {

        return textBlock.getAtomicTextBlocks().stream().map(AtomicTextBlock::getId).toArray(Long[]::new);
    }


    private DocumentPage toPageData(Page p) {

        return DocumentPage.builder().rotation(p.getRotation()).height(p.getHeight()).width(p.getWidth()).number(p.getNumber()).build();
    }


    private DocumentTextData toAtomicTextBlockData(AtomicTextBlock atomicTextBlock) {

        return DocumentTextData.builder()
                .id(atomicTextBlock.getId())
                .page(atomicTextBlock.getPage().getNumber().longValue())
                .searchText(atomicTextBlock.getSearchText())
                .numberOnPage(atomicTextBlock.getNumberOnPage())
                .start(atomicTextBlock.getBoundary().start())
                .end(atomicTextBlock.getBoundary().end())
                .lineBreaks(toPrimitiveIntArray(atomicTextBlock.getLineBreaks()))
                .build();
    }


    private DocumentPositionData toAtomicPositionBlockData(AtomicTextBlock atomicTextBlock) {

        return DocumentPositionData.builder()
                .id(atomicTextBlock.getId())
                .positions(toPrimitiveFloatMatrix(atomicTextBlock.getPositions()))
                .stringIdxToPositionIdx(toPrimitiveIntArray(atomicTextBlock.getStringIdxToPositionIdx()))
                .build();
    }


    private static float[][] toPrimitiveFloatMatrix(List<Rectangle2D> positions) {

        float[][] positionMatrix = new float[positions.size()][];
        for (int i = 0; i < positions.size(); i++) {
            positionMatrix[i] = toArray(positions.get(i));
        }
        return positionMatrix;
    }


    private static float[] toArray(Rectangle2D positions) {

        return new float[]{(float) positions.getMinX(), (float) positions.getMinY(), (float) positions.getWidth(), (float) positions.getHeight()};
    }


    private int[] toPrimitiveIntArray(List<Integer> list) {

        return list.stream().mapToInt(Integer::intValue).toArray();
    }

}
