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

package documenttext

import (
	"context"
	"encoding/json"
	"log/slog"
	"strings"

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

type cell struct {
	value      string
	row        int
	column     int
	isSelected bool
}

func (s *Service) addTable(ctx context.Context, id uuid.UUID, blockMap map[uuid.UUID]types.Block, builder *strings.Builder) error {
	params := tableParams{
		maxRow:    0,
		maxColumn: 0,
		table:     []cell{},
	}

	for _, uid := range s.GetDirectChildren(id, blockMap) {
		s.addComponentTable(ctx, uid, blockMap, &params)
	}

	invalidRows := make([]bool, params.maxRow+1)
	for _, cell := range params.table {
		if !cell.isSelected {
			invalidRows[cell.row] = true
		}
	}

	tableArr := make([][]string, params.maxRow+1)
	for i := range tableArr {
		tableArr[i] = make([]string, params.maxColumn+1)
	}

	for _, cell := range params.table {
		tableArr[cell.row][cell.column] = cell.value
	}

	finalTable := [][]string{}
	for i, row := range tableArr {
		if !invalidRows[i] {
			finalTable = append(finalTable, row)
		}
	}

	jsonData, err := json.Marshal(finalTable)
	if err != nil {
		return err
	}

	builder.WriteString("-------Table Start--------\n")
	builder.WriteString(string(jsonData))
	builder.WriteString("\n-------Table End--------\n")

	return nil
}

type tableParams struct {
	maxRow    int
	maxColumn int
	table     []cell
}

func (s *Service) addComponentTable(ctx context.Context, id uuid.UUID, blockMap map[uuid.UUID]types.Block, params *tableParams) {
	block := blockMap[id]

	switch block.BlockType {
	case types.BlockTypeMergedCell:
		content := ""
		addRow := true
		for _, uid := range s.getRelationshipIds(block) {
			if blockMap[uid].BlockType != types.BlockTypeCell {
				slog.Warn("not a cell", "block_type", blockMap[uid].BlockType)
				continue
			}
			content, addRow = s.getCellContent(uid, blockMap)
			if content != "" {
				break
			}
		}
		startColumnIndex := int(*block.ColumnIndex - 1)
		startRowIndex := int(*block.RowIndex - 1)
		endColumnIndex := startColumnIndex + int(*block.ColumnSpan) - 1
		endRowIndex := startRowIndex + int(*block.RowSpan) - 1
		for row := startRowIndex; row <= endRowIndex; row++ {
			for column := startColumnIndex; column <= endColumnIndex; column++ {
				s.addCellToTable(
					row,
					column,
					content,
					addRow,
					params,
				)
			}
		}
	case types.BlockTypeCell:
		content, addRow := s.getCellContent(id, blockMap)
		s.addCellToTable(
			int(*block.RowIndex)-1,
			int(*block.ColumnIndex)-1,
			content,
			addRow,
			params,
		)
	default:
		slog.Info("unhandled block type in table", "type", block.BlockType)
	}
}

func (s *Service) getCellContent(id uuid.UUID, blockMap map[uuid.UUID]types.Block) (string, bool) {
	var str strings.Builder
	var isSelected *bool
	for _, uid := range s.getRelationshipIds(blockMap[id]) {
		if blockMap[uid].BlockType == types.BlockTypeWord {
			if str.String() != "" {
				str.WriteRune(' ')
			}
			str.WriteString(*blockMap[uid].Text)
		}

		if blockMap[uid].BlockType != types.BlockTypeSelectionElement {
			continue
		}

		if isSelected == nil || !*isSelected {
			selected := true
			if blockMap[uid].SelectionStatus == types.SelectionStatusNotSelected {
				selected = false
			}
			isSelected = &selected
		}
	}

	selected := true
	if isSelected != nil {
		selected = *isSelected
	}

	return str.String(), selected
}

func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string, isSelected bool, params *tableParams) {
	if params.maxRow < rowIndex {
		params.maxRow = rowIndex
	}
	if params.maxColumn < columnIndex {
		params.maxColumn = columnIndex
	}
	params.table = append(params.table, cell{
		column:     columnIndex,
		row:        rowIndex,
		value:      content,
		isSelected: isSelected,
	})
}