query-orchestration/internal/document/text/getPages.go

package documenttext

import (
	"context"
	"fmt"
	"log/slog"
	"strings"

	documenttypes "queryorchestration/internal/document/types"

	"github.com/aws/aws-sdk-go-v2/service/textract"
	"github.com/aws/aws-sdk-go-v2/service/textract/types"
	"github.com/google/uuid"
)

func (s *Service) getPageText(ctx context.Context, document documenttypes.File, index int) (string, error) {
	res, err := s.GetBasePage(ctx, document, index)
	if err != nil {
		return "", err
	}

	var pageBlock *types.Block
	for _, block := range res.Blocks {
		if block.BlockType == types.BlockTypePage {
			pageBlock = &block
			break
		}
	}
	if pageBlock == nil {
		return "", fmt.Errorf("no page block found for index: %d", index)
	}

	blockMap := s.getBlockMap(res.Blocks)

	lines, tables, tableBlockMap, found, err := s.getPageElements(ctx, document, index, *pageBlock, blockMap)
	if err != nil {
		return "", err
	}

	return s.buildPageText(ctx, lines, document, blockMap, index, tables, tableBlockMap, found)
}

func (s *Service) getPageElements(ctx context.Context, document documenttypes.File, index int, pageBlock types.Block, blockMap map[uuid.UUID]types.Block) ([]uuid.UUID, []uuid.UUID, map[uuid.UUID]types.Block, map[uuid.UUID]*bool, error) {
	lines := []uuid.UUID{}
	hasTable := false
	found := map[uuid.UUID]*bool{}

	for _, id := range s.getChildIds(pageBlock) {
		if found[id] != nil && !*found[id] {
			continue
		}

		if s.isLayout(blockMap[id].BlockType) {
			lines = append(lines, id)
			f := true
			found[id] = &f
		}
		if blockMap[id].BlockType == types.BlockTypeLayoutTable {
			hasTable = true
		}
		if blockMap[id].BlockType == types.BlockTypeLayoutTable ||
			blockMap[id].BlockType == types.BlockTypeLayoutList {
			for _, childId := range s.getChildIds(blockMap[id]) {
				f := false
				found[childId] = &f
			}
		}
	}

	var tables []uuid.UUID
	var tableBlockMap map[uuid.UUID]types.Block
	if hasTable {
		ids, blockMap, err := s.getTables(ctx, document, index)
		if err != nil {
			return nil, nil, nil, nil, err
		}
		tables = ids
		tableBlockMap = blockMap
	}

	return lines, tables, tableBlockMap, found, nil
}

func (s *Service) isLayout(blockType types.BlockType) bool {
	return blockType == types.BlockTypeLayoutText ||
		blockType == types.BlockTypeLayoutTitle ||
		blockType == types.BlockTypeLayoutHeader ||
		blockType == types.BlockTypeLayoutFooter ||
		blockType == types.BlockTypeLayoutSectionHeader ||
		blockType == types.BlockTypeLayoutPageNumber ||
		blockType == types.BlockTypeLayoutList ||
		blockType == types.BlockTypeLayoutFigure ||
		blockType == types.BlockTypeLayoutTable ||
		blockType == types.BlockTypeLayoutKeyValue ||
		blockType == types.BlockTypeSignature
}

type buildParams struct {
	numOfSignatures int
	currentTable    int
}

func (s *Service) buildPageText(ctx context.Context, lines []uuid.UUID, document documenttypes.File, blockMap map[uuid.UUID]types.Block, index int, tables []uuid.UUID, tableBlockMap map[uuid.UUID]types.Block, found map[uuid.UUID]*bool) (string, error) {
	var builder strings.Builder
	params := buildParams{
		numOfSignatures: 0,
		currentTable:    0,
	}

	for _, id := range lines {
		if found[id] == nil || !*found[id] {
			continue
		}
		err := s.addComponentText(ctx, blockMap[id], blockMap, tables, tableBlockMap, &params, &builder)
		if err != nil {
			return "", err
		}
	}

	if params.numOfSignatures > 0 {
		builder.WriteString(fmt.Sprintf("\nThis page has %d signature.\n", params.numOfSignatures))
	}

	return builder.String(), nil
}

func (s *Service) addComponentText(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, tables []uuid.UUID, tableBlockMap map[uuid.UUID]types.Block, params *buildParams, builder *strings.Builder) error {
	switch block.BlockType {
	case types.BlockTypeLayoutText, types.BlockTypeLayoutTitle, types.BlockTypeLayoutHeader, types.BlockTypeLayoutFooter, types.BlockTypeLayoutSectionHeader, types.BlockTypeLayoutPageNumber, types.BlockTypeLayoutList, types.BlockTypeLayoutFigure, types.BlockTypeLayoutKeyValue:
		for _, id := range s.getChildIds(block) {
			err := s.addComponentText(ctx, blockMap[id], blockMap, tables, tableBlockMap, params, builder)
			if err != nil {
				return err
			}
		}
	case types.BlockTypeLine:
		_, err := builder.WriteString(fmt.Sprintf("%s\n", *block.Text))
		if err != nil {
			slog.Error("unable to write string", "error", err)
			return err
		}
	case types.BlockTypeWord:
		_, err := builder.WriteString(fmt.Sprintf(" %s", *block.Text))
		if err != nil {
			slog.Error("unable to write string", "error", err)
			return err
		}
	case types.BlockTypeSignature:
		params.numOfSignatures++
	case types.BlockTypeLayoutTable:
		tableId := tables[params.currentTable]
		params.currentTable++
		tableBlock := tableBlockMap[tableId]
		err := s.addTable(ctx, tableBlock, tableBlockMap, builder)
		if err != nil {
			return err
		}
	}

	return nil
}

func (s *Service) GetBasePage(ctx context.Context, doc documenttypes.File, index int) (textract.AnalyzeDocumentOutput, error) {
	pageBytes, err := doc.GetPageAsPNG(ctx, index)
	if err != nil {
		return textract.AnalyzeDocumentOutput{}, err
	}

	res, err := s.cfg.GetTextractClient().AnalyzeDocument(ctx, &textract.AnalyzeDocumentInput{
		Document: &types.Document{
			Bytes: pageBytes,
		},
		FeatureTypes: []types.FeatureType{
			types.FeatureTypeLayout,
			types.FeatureTypeSignatures,
		},
	})
	if err != nil {
		return textract.AnalyzeDocumentOutput{}, err
	}

	return *res, err
}

func (s *Service) getBlockMap(blocksList []types.Block) map[uuid.UUID]types.Block {
	blocks := make(map[uuid.UUID]types.Block, len(blocksList))
	for _, block := range blocksList {
		if block.Id == nil {
			continue
		}
		id, err := uuid.Parse(*block.Id)
		if err != nil {
			slog.Warn("error parsing block id", "error", err, "id", *block.Id)
			continue
		}

		blocks[id] = block
	}

	return blocks
}

func (s *Service) getChildIds(block types.Block) []uuid.UUID {
	childIds := []uuid.UUID{}
	for _, relationship := range block.Relationships {
		if relationship.Type != types.RelationshipTypeChild {
			slog.Info("not a child element", "type", relationship.Type)
			continue
		}

		for _, idStr := range relationship.Ids {
			id, err := uuid.Parse(idStr)
			if err != nil {
				slog.Error("invalid child uuid", "id", id, "error", err)
				continue
			}

			childIds = append(childIds, id)
		}
	}

	return childIds
}