package documenttext

import (
	"context"
	"encoding/json"
	"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"
)

type cell struct {
	value  string
	row    int
	column int
}

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

	for _, uid := range s.getChildIds(block) {
		childBlock := blockMap[uid]
		if childBlock.BlockType != types.BlockTypeCell {
			continue
		}

		var str strings.Builder
		for _, uid := range s.getChildIds(childBlock) {
			if blockMap[uid].BlockType == types.BlockTypeWord {
				if str.String() != "" {
					str.WriteString(" ")
				}
				str.WriteString(*blockMap[uid].Text)
			}
			if maxRow < int(*childBlock.RowIndex) {
				maxRow = int(*childBlock.RowIndex)
			}
			if maxColumn < int(*childBlock.ColumnIndex) {
				maxColumn = int(*childBlock.ColumnIndex)
			}
			table = append(table, cell{
				column: int(*childBlock.ColumnIndex) - 1,
				row:    int(*childBlock.RowIndex) - 1,
				value:  str.String(),
			})
		}
	}

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

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

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

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

	return nil
}

func (s *Service) getTables(ctx context.Context, document documenttypes.File, index int) ([]uuid.UUID, map[uuid.UUID]types.Block, error) {
	pageBytes, err := document.GetPageAsPNG(ctx, index)
	if err != nil {
		return nil, nil, err
	}

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

	tables := []uuid.UUID{}
	for _, block := range res.Blocks {
		if block.BlockType == types.BlockTypeTable {
			id, err := uuid.Parse(*block.Id)
			if err != nil {
				slog.Warn("error parsing block id", "error", err, "id", *block.Id)
				continue
			}

			tables = append(tables, id)
		}
	}
	blocks := s.getBlockMap(res.Blocks)

	return tables, blocks, nil
}