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

import java.awt.geom.Rectangle2D;
import java.util.List;
import java.util.Locale;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

import com.knecon.fforesight.service.layoutparser.internal.api.data.redaction.NodeType;
import com.knecon.fforesight.service.layoutparser.internal.api.data.taas.ParagraphData;
import com.knecon.fforesight.service.layoutparser.internal.api.data.taas.Range;
import com.knecon.fforesight.service.layoutparser.internal.api.data.taas.ResearchDocumentData;
import com.knecon.fforesight.service.layoutparser.internal.api.data.taas.RowData;
import com.knecon.fforesight.service.layoutparser.internal.api.data.taas.StructureObject;
import com.knecon.fforesight.service.layoutparser.internal.api.data.taas.TableData;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.Boundary;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Document;
import com.knecon.fforesight.service.layoutparser.processor.model.graph.nodes.Page;
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.TextBlock;

public class TaasDocumentDataMapper {

    public static ResearchDocumentData fromDocument(Document document) {

        AtomicInteger structureObjectNumber = new AtomicInteger();
        List<StructureObject> structureObjects = document.streamAllSubNodes()
                .filter(node -> !node.getType().equals(NodeType.TABLE_CELL))
                .filter(node -> !node.getType().equals(NodeType.SECTION))
                .map(node -> {
                    if (node.getType().equals(NodeType.TABLE)) {
                        return TaasDocumentDataMapper.fromTableWithTableData((Table) node, structureObjectNumber.getAndIncrement());
                    } else {
                        return TaasDocumentDataMapper.fromSemanticNodeWithParagraphData(node, structureObjectNumber.getAndIncrement());
                    }
                })
                .toList();
        return ResearchDocumentData.builder().structureObjects(structureObjects).build();
    }


    public static ParagraphData fromTextBlock(String classification, TextBlock textBlock) {

        return ParagraphData.builder()
                .boldTextBoundaries(textBlock.getBoldTextBoundaries().stream().map(b -> new Range(b.start(), b.end())).toList())
                .italicTextBoundaries(textBlock.getItalicTextBoundaries().stream().map(b -> new Range(b.start(), b.end())).toList())
                .text(textBlock.getSearchText())
                .linebreaks(textBlock.getLineBreaks())
                .classification(classification)
                .orientation(textBlock.getOrientation())
                .textDirection(textBlock.getTextDirection())
                .build();
    }


    public static TableData fromTable(Table table) {

        List<RowData> rowData = IntStream.range(0, table.getNumberOfRows())
                .boxed()
                .map(rowIdx -> table.streamRow(rowIdx).toList())
                .map(TaasDocumentDataMapper::fromTableCells)
                .toList();
        return new TableData(rowData, table.getNumberOfCols(), table.getNumberOfRows());
    }


    public static RowData fromTableCells(List<TableCell> tableCells) {

        if (tableCells.isEmpty()) {
            throw new IllegalArgumentException("no table cells provided");
        }
        boolean header = tableCells.stream().allMatch(TableCell::isHeader);
        Page firstPage = tableCells.get(0).getFirstPage();
        Rectangle2D bBox = tableCells.stream().map(TableCell::getBBox).reduce((map1, map2) -> {
            map1.forEach((page, rectangle) -> map2.merge(page, rectangle, (rect1, rect2) -> rect1.createUnion(rect2).getBounds2D()));
            return map2;
        }).orElseThrow().get(firstPage);
        List<TextBlock> textBlocks = tableCells.stream().map(TableCell::getTextBlock).toList();
        return new RowData(header, textBlocks.stream().map(textBlock -> TaasDocumentDataMapper.fromTextBlock("table_cell", textBlock)).toList(), toFloatArray(bBox));
    }


    private static Range toRange(Boundary boundary) {

        return new Range(boundary.start(), boundary.end());
    }


    private static List<Range> toRange(List<Boundary> boundary) {

        return boundary.stream().map(TaasDocumentDataMapper::toRange).toList();
    }


    public static StructureObject fromSemanticNodeWithParagraphData(SemanticNode semanticNode, Integer structureObjectNumber) {

        Page page = semanticNode.getFirstPage();
        Rectangle2D bBox = semanticNode.getBBox().get(page);
        return StructureObject.builder()
                .structureObjectNumber(structureObjectNumber)
                .boundingBox(toFloatArray(bBox))
                .stringOffset(semanticNode.getBoundary().start())
                .page(page.getNumber())
                .paragraph(TaasDocumentDataMapper.fromTextBlock(semanticNode.getType().toString().toLowerCase(Locale.ROOT), semanticNode.getTextBlock()))
                .table(null)
                .build();
    }


    public static StructureObject fromTableWithTableData(Table table, int structureObjectNumber) {

        Page page = table.getFirstPage();
        Rectangle2D bBox = table.getBBox().get(page);
        return StructureObject.builder()
                .structureObjectNumber(structureObjectNumber)
                .boundingBox(toFloatArray(bBox))
                .stringOffset(table.getBoundary().start())
                .page(page.getNumber())
                .paragraph(null)
                .table(TaasDocumentDataMapper.fromTable(table))
                .build();
    }


    private static float[] toFloatArray(Rectangle2D bBox) {

        return new float[]{(float) bBox.getX(), (float) bBox.getY(), (float) bBox.getWidth(), (float) bBox.getHeight()};
    }

}
