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

import static com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentTextDataProto.AllDocumentTextData;
import static com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentTextDataProto.DocumentTextData;

import java.awt.geom.Rectangle2D;
import java.util.ArrayList;
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.DocumentPageProto.AllDocumentPages;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentPageProto.DocumentPage;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentPositionDataProto.AllDocumentPositionData;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentPositionDataProto.DocumentPositionData;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentPositionDataProto.DocumentPositionData.Position;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentStructureProto.DocumentStructure;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.DocumentStructureWrapper;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.EntryDataProto.EntryData;
import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.LayoutEngineProto.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();

        AllDocumentTextData allDocumentTextData = AllDocumentTextData.newBuilder().addAllDocumentTextData(documentTextData).build();

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

        AllDocumentPositionData allDocumentPositionData = AllDocumentPositionData.newBuilder().addAllDocumentPositionData(atomicPositionBlockData).build();

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

        List<DocumentPage> documentPageData = document.getPages()
                .stream()
                .map(DocumentDataMapper::toPageData)
                .toList();

        AllDocumentPages allDocumentPages = AllDocumentPages.newBuilder().addAllDocumentPages(documentPageData).build();
        DocumentStructureWrapper tableOfContentsData = toDocumentTreeData(document.getDocumentTree());
        return DocumentData.builder()
                .documentTextData(allDocumentTextData)
                .documentPositions(allDocumentPositionData)
                .documentPages(allDocumentPages)
                .documentStructureWrapper(tableOfContentsData)
                .build();
    }


    private DocumentStructureWrapper toDocumentTreeData(DocumentTree documentTree) {

        return new DocumentStructureWrapper(DocumentStructure.newBuilder().setRoot(toEntryData(documentTree.getRoot())).build());
    }


    private EntryData toEntryData(DocumentTree.Entry entry) {

        List<Long> atomicTextBlocks;

        if (entry.getNode().isLeaf()) {
            atomicTextBlocks = toAtomicTextBlockIds(entry.getNode().getLeafTextBlock());
        } else {
            atomicTextBlocks = new ArrayList<>();
        }

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

        var documentBuilder = EntryData.newBuilder()
                .addAllTreeId(entry.getTreeId())
                .addAllChildren(entry.getChildren()
                                        .stream()
                                        .map(DocumentDataMapper::toEntryData)
                                        .toList())
                .setType(entry.getType())
                .addAllAtomicBlockIds(atomicTextBlocks)
                .addAllPageNumbers(entry.getNode().getPages()
                                           .stream()
                                           .map(Page::getNumber)
                                           .map(Integer::longValue)
                                           .toList())
                .putAllProperties(properties);
        if (entry.getNode() != null) {
            documentBuilder.addAllEngines(entry.getNode().getEngines());
        } else {
            documentBuilder.addAllEngines(new HashSet<>(Set.of(LayoutEngine.ALGORITHM)));
        }
        return documentBuilder.build();
    }


    private List<Long> toAtomicTextBlockIds(TextBlock textBlock) {

        return textBlock.getAtomicTextBlocks()
                .stream()
                .map(AtomicTextBlock::getId)
                .toList();
    }


    private DocumentPage toPageData(Page p) {

        return DocumentPage.newBuilder().setRotation(p.getRotation()).setHeight(p.getHeight()).setWidth(p.getWidth()).setNumber(p.getNumber()).build();
    }


    private DocumentTextData toAtomicTextBlockData(AtomicTextBlock atomicTextBlock) {

        return DocumentTextData.newBuilder()
                .setId(atomicTextBlock.getId())
                .setPage(atomicTextBlock.getPage().getNumber().longValue())
                .setSearchText(atomicTextBlock.getSearchText())
                .setNumberOnPage(atomicTextBlock.getNumberOnPage())
                .setStart(atomicTextBlock.getTextRange().start())
                .setEnd(atomicTextBlock.getTextRange().end())
                .addAllLineBreaks(atomicTextBlock.getLineBreaks())
                .build();
    }


    private DocumentPositionData toAtomicPositionBlockData(AtomicTextBlock atomicTextBlock) {

        return DocumentPositionData.newBuilder()
                .setId(atomicTextBlock.getId())
                .addAllPositions(toPositions(atomicTextBlock.getPositions()))
                .addAllStringIdxToPositionIdx(atomicTextBlock.getStringIdxToPositionIdx())
                .build();
    }


    private static List<Position> toPositions(List<Rectangle2D> rects) {

        List<Position> positions = new ArrayList<>();
        for (Rectangle2D rect : rects) {
            positions.add(toPosition(rect));
        }
        return positions;
    }


    private static Position toPosition(Rectangle2D rect) {

        return Position.newBuilder().addValue((float) rect.getMinX()).addValue((float) rect.getMinY()).addValue((float) rect.getWidth()).addValue((float) rect.getHeight()).build();
    }

}
