Laniakea/tgapi/rich_utils.go

package tgapi

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
)

const (
	maximumRichJSONDepth = 64
	maximumRichJSONNodes = 10_000
)

var errInvalidUnknownRichJSON = errors.New("invalid JSON in unknown rich object")

var knownRichTextTypes = map[string]bool{
	"url": true, "email_address": true, "phone_number": true,
	"bank_card_number": true, "mention": true, "hashtag": true,
	"cashtag": true, "bot_command": true, "anchor_link": true,
	"reference": true, "reference_link": true, "date_time": true,
	"text_mention": true, "custom_emoji": true,
	"mathematical_expression": true, "anchor": true, "button": true,
}

// UnmarshalRichText parses a RichText tree from JSON: a string, an array, or
// a typed object. Unknown object types are preserved as RichTextUnknown without
// discarding fields.
//
// Since: Bot API 10.1
func UnmarshalRichText(data []byte) (RichText, error) {
	if err := validateRichJSON(data); err != nil {
		return nil, err
	}
	return unmarshalRichText(data)
}

func unmarshalRichText(data []byte) (RichText, error) {
	if bytes.Equal(bytes.TrimSpace(data), []byte("null")) {
		return nil, fmt.Errorf("richtext: null is not a rich text value")
	}
	// 1. string
	var s string
	if err := json.Unmarshal(data, &s); err == nil {
		return RichTextPlain(s), nil
	}
	// 2. array
	var raw []json.RawMessage
	if err := json.Unmarshal(data, &raw); err == nil {
		arr := make(RichTextArray, len(raw))
		for i, it := range raw {
			rt, err := unmarshalRichText(it)
			if err != nil {
				return nil, err
			}
			arr[i] = rt
		}
		return arr, nil
	}
	// 3. object -> dispatch on type, grabbing the raw text along the way
	var head struct {
		Type string          `json:"type"`
		Text json.RawMessage `json:"text"`
	}
	if err := json.Unmarshal(data, &head); err != nil {
		return nil, fmt.Errorf("richtext: not a string, array or object: %w", err)
	}
	if head.Type == "" {
		return nil, errors.New("richtext: object type is required")
	}
	if !richTextWrapTags[head.Type] && !knownRichTextTypes[head.Type] {
		return RichTextUnknown{Raw: cloneRawJSON(data)}, nil
	}

	// Recursively parse the nested text, if any.
	var inner RichText
	if len(head.Text) > 0 {
		var err error
		if inner, err = unmarshalRichText(head.Text); err != nil {
			return nil, fmt.Errorf("richtext %q: bad text: %w", head.Type, err)
		}
	}

	if richTextWrapTags[head.Type] {
		return RichTextWrap{Tag: head.Type, Text: inner}, nil
	}

	switch head.Type {
	case "url":
		var v struct {
			URL string `json:"url"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextURL{inner, v.URL}, nil
	case "email_address":
		var v struct {
			V string `json:"email_address"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextEmailAddress{inner, v.V}, nil
	case "phone_number":
		var v struct {
			V string `json:"phone_number"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextPhoneNumber{inner, v.V}, nil
	case "bank_card_number":
		var v struct {
			V string `json:"bank_card_number"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextBankCardNumber{inner, v.V}, nil
	case "mention":
		var v struct {
			V string `json:"username"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextMention{inner, v.V}, nil
	case "hashtag":
		var v struct {
			V string `json:"hashtag"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextHashtag{inner, v.V}, nil
	case "cashtag":
		var v struct {
			V string `json:"cashtag"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextCashtag{inner, v.V}, nil
	case "bot_command":
		var v struct {
			V string `json:"bot_command"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextBotCommand{inner, v.V}, nil
	case "anchor_link":
		var v struct {
			V string `json:"anchor_name"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextAnchorLink{inner, v.V}, nil
	case "reference":
		var v struct {
			V string `json:"name"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextReference{inner, v.V}, nil
	case "reference_link":
		var v struct {
			V string `json:"reference_name"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextReferenceLink{inner, v.V}, nil
	case "date_time":
		var v struct {
			UnixTime       int64  `json:"unix_time"`
			DateTimeFormat string `json:"date_time_format"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextDateTime{inner, v.UnixTime, v.DateTimeFormat}, nil
	case "text_mention":
		var v struct {
			User User `json:"user"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextTextMention{inner, v.User}, nil

	// --- leaves without text ---
	case "custom_emoji":
		var v struct {
			ID  string `json:"custom_emoji_id"`
			Alt string `json:"alternative_text"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextCustomEmoji{v.ID, v.Alt}, nil
	case "mathematical_expression":
		var v struct {
			Expression string `json:"expression"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextMathematicalExpression{v.Expression}, nil
	case "anchor":
		var v struct {
			Name string `json:"name"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichTextAnchor{v.Name}, nil
	case "button":
		var raw struct {
			Button RichMessageButton `json:"button"`
		}
		if err := json.Unmarshal(data, &raw); err != nil {
			return nil, err
		}
		return RichTextButton{Button: raw.Button}, nil

	default:
		return nil, fmt.Errorf("richtext: unsupported type %q", head.Type)
	}
}

// UnmarshalRichBlock parses a single RichBlock from JSON, dispatching on the
// type tag. Unknown object types are preserved as RichBlockUnknown without
// discarding fields.
//
// Since: Bot API 10.1
func UnmarshalRichBlock(data []byte) (RichBlock, error) {
	if err := validateRichJSON(data); err != nil {
		return nil, err
	}
	return unmarshalRichBlock(data)
}

func unmarshalRichBlock(data []byte) (RichBlock, error) {
	var head struct {
		Type string          `json:"type"`
		Text json.RawMessage `json:"text"`
	}
	if err := json.Unmarshal(data, &head); err != nil {
		return nil, fmt.Errorf("richblock: %w", err)
	}
	if head.Type == "" {
		return nil, errors.New("richblock: object type is required")
	}

	if richBlockWrapTags[head.Type] {
		text, err := parseOptRichText(head.Text)
		if err != nil {
			return nil, fmt.Errorf("richblock %q: text: %w", head.Type, err)
		}
		return RichBlockWrap{Tag: head.Type, Text: text}, nil
	}

	switch head.Type {
	case "heading":
		var v struct {
			Size int `json:"size"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		text, err := parseOptRichText(head.Text)
		if err != nil {
			return nil, fmt.Errorf("richblock %q: text: %w", head.Type, err)
		}
		return RichBlockSectionHeading{text, v.Size}, nil

	case "pre":
		var v struct {
			Language string `json:"language"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		text, err := parseOptRichText(head.Text)
		if err != nil {
			return nil, fmt.Errorf("richblock %q: text: %w", head.Type, err)
		}
		return RichBlockPreformatted{text, v.Language}, nil

	case "blockquote":
		var raw struct {
			Blocks json.RawMessage `json:"blocks"`
			Credit json.RawMessage `json:"credit"`
		}
		if err := json.Unmarshal(data, &raw); err != nil {
			return nil, err
		}
		blocks, err := unmarshalRichBlocks(raw.Blocks)
		if err != nil {
			return nil, err
		}
		credit, err := parseOptRichText(raw.Credit)
		if err != nil {
			return nil, fmt.Errorf("richblock %q: credit: %w", head.Type, err)
		}
		return RichBlockQuotation{blocks, credit}, nil

	case "expandable_blockquote", "pullquote":
		var raw struct {
			Credit json.RawMessage `json:"credit"`
		}
		if err := json.Unmarshal(data, &raw); err != nil {
			return nil, err
		}
		text, err := parseOptRichText(head.Text)
		if err != nil {
			return nil, fmt.Errorf("richblock %q: text: %w", head.Type, err)
		}
		credit, err := parseOptRichText(raw.Credit)
		if err != nil {
			return nil, fmt.Errorf("richblock %q: credit: %w", head.Type, err)
		}
		if head.Type == "expandable_blockquote" {
			return RichBlockExpandableBlockQuotation{text, credit}, nil
		}
		return RichBlockPullQuotation{text, credit}, nil

	case "list":
		var v struct {
			Items []RichBlockListItem `json:"items"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockList{v.Items}, nil

	case "collage":
		var raw struct {
			Blocks  json.RawMessage   `json:"blocks"`
			Caption *RichBlockCaption `json:"caption"`
		}
		if err := json.Unmarshal(data, &raw); err != nil {
			return nil, err
		}
		blocks, err := unmarshalRichBlocks(raw.Blocks)
		if err != nil {
			return nil, err
		}
		return RichBlockCollage{blocks, raw.Caption}, nil

	case "slideshow":
		var raw struct {
			Blocks  json.RawMessage   `json:"blocks"`
			Caption *RichBlockCaption `json:"caption"`
		}
		if err := json.Unmarshal(data, &raw); err != nil {
			return nil, err
		}
		blocks, err := unmarshalRichBlocks(raw.Blocks)
		if err != nil {
			return nil, err
		}
		return RichBlockSlideshow{blocks, raw.Caption}, nil

	case "details":
		var raw struct {
			Summary json.RawMessage `json:"summary"`
			Blocks  json.RawMessage `json:"blocks"`
			IsOpen  bool            `json:"is_open"`
		}
		if err := json.Unmarshal(data, &raw); err != nil {
			return nil, err
		}
		summary, err := parseOptRichText(raw.Summary)
		if err != nil {
			return nil, fmt.Errorf("richblock %q: summary: %w", head.Type, err)
		}
		blocks, err := unmarshalRichBlocks(raw.Blocks)
		if err != nil {
			return nil, err
		}
		return RichBlockDetails{summary, blocks, raw.IsOpen}, nil

	case "table":
		var raw struct {
			Cells      [][]RichBlockTableCell `json:"cells"`
			IsBordered bool                   `json:"is_bordered"`
			IsStriped  bool                   `json:"is_striped"`
			IsCompact  bool                   `json:"is_compact"`
			Caption    json.RawMessage        `json:"caption"`
		}
		if err := json.Unmarshal(data, &raw); err != nil {
			return nil, err
		}
		caption, err := parseOptRichText(raw.Caption)
		if err != nil {
			return nil, fmt.Errorf("richblock %q: caption: %w", head.Type, err)
		}
		return RichBlockTable{raw.Cells, raw.IsBordered, raw.IsStriped, raw.IsCompact, caption}, nil

	case "map":
		var v struct {
			Location Location          `json:"location"`
			Zoom     int               `json:"zoom"`
			Width    int               `json:"width"`
			Height   int               `json:"height"`
			Caption  *RichBlockCaption `json:"caption"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockMap{v.Location, v.Zoom, v.Width, v.Height, v.Caption}, nil

	case "buttons":
		var block RichBlockButtons
		if err := json.Unmarshal(data, &block); err != nil {
			return nil, err
		}
		return block, nil
	case "document":
		var block RichBlockDocument
		if err := json.Unmarshal(data, &block); err != nil {
			return nil, err
		}
		return block, nil
	case "photo":
		var v struct {
			Photo      []PhotoSize       `json:"photo"`
			HasSpoiler bool              `json:"has_spoiler"`
			Caption    *RichBlockCaption `json:"caption"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockPhoto{v.Photo, v.HasSpoiler, v.Caption}, nil

	case "video":
		var v struct {
			Video      Video             `json:"video"`
			HasSpoiler bool              `json:"has_spoiler"`
			Caption    *RichBlockCaption `json:"caption"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockVideo{v.Video, v.HasSpoiler, v.Caption}, nil

	case "audio":
		var v struct {
			Audio   Audio             `json:"audio"`
			Caption *RichBlockCaption `json:"caption"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockAudio{v.Audio, v.Caption}, nil

	case "animation":
		var v struct {
			Animation  Animation         `json:"animation"`
			HasSpoiler bool              `json:"has_spoiler"`
			Caption    *RichBlockCaption `json:"caption"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockAnimation{v.Animation, v.HasSpoiler, v.Caption}, nil

	case "voice_note":
		var v struct {
			VoiceNote Voice             `json:"voice_note"`
			Caption   *RichBlockCaption `json:"caption"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockVoiceNote{v.VoiceNote, v.Caption}, nil

	case "divider":
		return RichBlockDivider{}, nil

	case "mathematical_expression":
		var v struct {
			Expression string `json:"expression"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockMathematicalExpression{v.Expression}, nil

	case "anchor":
		var v struct {
			Name string `json:"name"`
		}
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, err
		}
		return RichBlockAnchor{v.Name}, nil

	default:
		return RichBlockUnknown{Raw: cloneRawJSON(data)}, nil
	}
}

// UnmarshalRichMessage parses a RichMessage and requires a non-null blocks array.
//
// Since: Bot API 10.1
func UnmarshalRichMessage(data []byte) (RichMessage, error) {
	return unmarshalRichMessageStrict(data)
}

func unmarshalRichMessage(data []byte) (RichMessage, error) {
	var raw struct {
		Blocks json.RawMessage `json:"blocks"`
		IsRTL  bool            `json:"is_rtl"`
	}
	if err := json.Unmarshal(data, &raw); err != nil {
		return RichMessage{}, fmt.Errorf("richmessage: %w", err)
	}
	blocks, err := unmarshalRichBlocks(raw.Blocks)
	if err != nil {
		return RichMessage{}, err
	}
	return RichMessage{blocks, raw.IsRTL}, nil
}

// UnmarshalRichMessageStrict parses a RichMessage and requires a non-null blocks array.
//
// Since: Bot API 10.1
func UnmarshalRichMessageStrict(data []byte) (RichMessage, error) {
	return unmarshalRichMessageStrict(data)
}

func unmarshalRichMessageStrict(data []byte) (RichMessage, error) {
	if err := validateRichJSON(data); err != nil {
		return RichMessage{}, err
	}
	var root map[string]json.RawMessage
	if err := json.Unmarshal(data, &root); err != nil {
		return RichMessage{}, fmt.Errorf("richmessage: %w", err)
	}
	blocks, ok := root["blocks"]
	if !ok || bytes.Equal(bytes.TrimSpace(blocks), []byte("null")) {
		return RichMessage{}, errors.New("richmessage: blocks must be a non-null array")
	}
	var rawBlocks []json.RawMessage
	if err := json.Unmarshal(blocks, &rawBlocks); err != nil {
		return RichMessage{}, errors.New("richmessage: blocks must be an array")
	}
	return unmarshalRichMessage(data)
}

// UnmarshalJSON implements json.Unmarshaler.
//
// Since: Bot API 10.1
func (m *RichMessage) UnmarshalJSON(data []byte) error {
	parsed, err := UnmarshalRichMessage(data)
	if err != nil {
		return err
	}
	*m = parsed
	return nil
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

// Optional RichText fields treat absent and null values as nil.
func parseOptRichText(raw json.RawMessage) (RichText, error) {
	if len(raw) == 0 || string(raw) == "null" {
		return nil, nil
	}
	return unmarshalRichText(raw)
}

func cloneRawJSON(raw []byte) json.RawMessage {
	return append(json.RawMessage(nil), raw...)
}

func unmarshalRichBlocks(raw json.RawMessage) ([]RichBlock, error) {
	if len(raw) == 0 || string(raw) == "null" {
		return nil, nil
	}
	var raws []json.RawMessage
	if err := json.Unmarshal(raw, &raws); err != nil {
		return nil, err
	}
	blocks := make([]RichBlock, len(raws))
	for i, r := range raws {
		b, err := unmarshalRichBlock(r)
		if err != nil {
			return nil, err
		}
		blocks[i] = b
	}
	return blocks, nil
}

func validateRichJSON(data []byte) error {
	decoder := json.NewDecoder(bytes.NewReader(data))
	depth := 0
	nodes := 0
	for {
		token, err := decoder.Token()
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return err
		}
		nodes++
		if nodes > maximumRichJSONNodes {
			return fmt.Errorf("%w: maximum %d", ErrRichJSONNodes, maximumRichJSONNodes)
		}
		delim, ok := token.(json.Delim)
		if !ok {
			continue
		}
		switch delim {
		case '{', '[':
			depth++
			if depth > maximumRichJSONDepth {
				return fmt.Errorf("%w: maximum %d", ErrRichJSONDepth, maximumRichJSONDepth)
			}
		case '}', ']':
			depth--
		}
	}
}