tgex-backend/telegram_bot/ui/message_format.go

package ui

import (
	"fmt"
	"sort"
	"strings"

	"github.com/NicoNex/echotron/v3"
)

type htmlTag struct {
	open  string
	close string
}

var simpleHTMLTags = map[echotron.MessageEntityType]htmlTag{
	echotron.BoldEntity:          {open: "<b>", close: "</b>"},
	echotron.ItalicEntity:        {open: "<i>", close: "</i>"},
	echotron.UnderlineEntity:     {open: "<u>", close: "</u>"},
	echotron.StrikethroughEntity: {open: "<s>", close: "</s>"},
	echotron.CodeEntity:          {open: "<code>", close: "</code>"},
	echotron.PreEntity:           {open: "<pre>", close: "</pre>"},
}

// FormatMessageHTML converts Telegram entities to HTML and escapes the rest.
// Invite links are normalized into <a class="tg-link"> placeholders for previews.
func FormatMessageHTML(message *echotron.Message) string {
	if message == nil {
		return ""
	}

	// Telegram sends formatting as entities, so we rebuild the HTML from offsets.
	text := message.Text
	entities := message.Entities
	if text == "" {
		text = message.Caption
		entities = message.CaptionEntities
	}

	if text == "" {
		return ""
	}

	if len(entities) == 0 {
		return EscapeHTML(text)
	}

	type tagSpan struct {
		open  string
		close string
		start int
		end   int
		len   int
	}

	runes := []rune(text)
	positions := make([]int, len(runes)+1)
	utf16Count := 0
	for i, r := range runes {
		positions[i] = utf16Count
		if r > 0xFFFF {
			utf16Count += 2
		} else {
			utf16Count++
		}
	}
	positions[len(runes)] = utf16Count

	utf16ToRuneIndex := func(utf16Index int) (int, bool) {
		i := sort.Search(len(positions), func(i int) bool { return positions[i] >= utf16Index })
		if i < len(positions) && positions[i] == utf16Index {
			return i, true
		}
		return 0, false
	}

	toRuneRange := func(offset, length int) (int, int, bool) {
		start, ok := utf16ToRuneIndex(offset)
		if !ok {
			return 0, 0, false
		}
		end, ok := utf16ToRuneIndex(offset + length)
		if !ok || end > len(runes) || end < start {
			return 0, 0, false
		}
		return start, end, true
	}

	// URL ranges are used to prevent nested tags inside links.
	urlRanges := make([][2]int, 0, len(entities))

	for _, entity := range entities {
		if entity == nil {
			continue
		}
		if entity.Type != echotron.UrlEntity && entity.Type != echotron.TextLinkEntity {
			continue
		}
		start, end, ok := toRuneRange(entity.Offset, entity.Length)
		if !ok {
			continue
		}
		urlRanges = append(urlRanges, [2]int{start, end})
		if entity.Type == echotron.UrlEntity {
			_ = string(runes[start:end])
		}
	}

	isInsideURL := func(offset, length int) bool {
		for _, urlRange := range urlRanges {
			if offset > urlRange[0] && offset+length <= urlRange[1] {
				return true
			}
			if offset < urlRange[1] && offset+length > urlRange[0] && (offset != urlRange[0] || length != urlRange[1]-urlRange[0]) {
				return true
			}
		}
		return false
	}

	var spans []tagSpan
	for _, entity := range entities {
		if entity == nil {
			continue
		}
		start, end, ok := toRuneRange(entity.Offset, entity.Length)
		if !ok {
			continue
		}

		if entity.Type != echotron.UrlEntity && entity.Type != echotron.TextLinkEntity && isInsideURL(start, end-start) {
			continue
		}

		var openTag, closeTag string
		if tag, ok := simpleHTMLTags[entity.Type]; ok {
			openTag, closeTag = tag.open, tag.close
		} else if entity.Type == echotron.UrlEntity {
			entityText := string(runes[start:end])
			if isInviteLink(entityText) {
				openTag = `<a class="tg-link">`
				closeTag = "</a>"
			} else {
				openTag = fmt.Sprintf("<a href=\"%s\">", entityText)
				closeTag = "</a>"
			}
		} else if entity.Type == echotron.TextLinkEntity {
			if entity.URL != "" {
				if isInviteLink(entity.URL) {
					openTag = `<a class="tg-link">`
					closeTag = "</a>"
				} else {
					openTag = fmt.Sprintf("<a href=\"%s\">", entity.URL)
					closeTag = "</a>"
				}
			}
		}

		if openTag != "" && closeTag != "" {
			spans = append(spans, tagSpan{
				open:  openTag,
				close: closeTag,
				start: start,
				end:   end,
				len:   end - start,
			})
		}
	}

	opens := make(map[int][]tagSpan)
	closes := make(map[int][]tagSpan)
	for _, span := range spans {
		opens[span.start] = append(opens[span.start], span)
		closes[span.end] = append(closes[span.end], span)
	}

	// Build the final text with tags inserted and HTML escaped.
	var result strings.Builder
	for i := 0; i <= len(runes); i++ {
		if closing, ok := closes[i]; ok {
			sort.Slice(closing, func(a, b int) bool { return closing[a].len < closing[b].len })
			for _, span := range closing {
				result.WriteString(span.close)
			}
		}
		if opening, ok := opens[i]; ok {
			sort.Slice(opening, func(a, b int) bool { return opening[a].len > opening[b].len })
			for _, span := range opening {
				result.WriteString(span.open)
			}
		}
		if i < len(runes) {
			switch runes[i] {
			case '<':
				result.WriteString("&lt;")
			case '>':
				result.WriteString("&gt;")
			case '&':
				result.WriteString("&amp;")
			default:
				result.WriteRune(runes[i])
			}
		}
	}

	return normalizeHTMLTags(result.String())
}

// isInviteLink detects Telegram invite links used in creatives.
func isInviteLink(url string) bool {
	return strings.Contains(url, "t.me/+") || strings.Contains(url, "t.me/joinchat/")
}

func normalizeHTMLTags(input string) string {
	if input == "" {
		return input
	}

	allowed := map[string]bool{
		"a":    true,
		"b":    true,
		"i":    true,
		"u":    true,
		"s":    true,
		"code": true,
		"pre":  true,
	}

	var out strings.Builder
	out.Grow(len(input))
	stack := make([]string, 0, 8)

	for i := 0; i < len(input); {
		if input[i] != '<' {
			out.WriteByte(input[i])
			i++
			continue
		}

		end := strings.IndexByte(input[i:], '>')
		if end == -1 {
			out.WriteByte(input[i])
			i++
			continue
		}

		end += i
		tag := input[i+1 : end]
		if tag == "" {
			out.WriteString(input[i : end+1])
			i = end + 1
			continue
		}

		isClosing := tag[0] == '/'
		tagName := tag
		if isClosing {
			tagName = tag[1:]
		}
		if space := strings.IndexByte(tagName, ' '); space != -1 {
			tagName = tagName[:space]
		}
		tagName = strings.TrimSpace(tagName)

		if !allowed[tagName] {
			out.WriteString(input[i : end+1])
			i = end + 1
			continue
		}

		if isClosing {
			if len(stack) > 0 && stack[len(stack)-1] == tagName {
				stack = stack[:len(stack)-1]
				out.WriteString(input[i : end+1])
			}
		} else {
			stack = append(stack, tagName)
			out.WriteString(input[i : end+1])
		}

		i = end + 1
	}

	for i := len(stack) - 1; i >= 0; i-- {
		out.WriteString("</")
		out.WriteString(stack[i])
		out.WriteString(">")
	}

	return out.String()
}

// EscapeHTML escapes HTML special characters for safe output.
func EscapeHTML(s string) string {
	s = strings.ReplaceAll(s, "&", "&amp;")
	s = strings.ReplaceAll(s, "<", "&lt;")
	s = strings.ReplaceAll(s, ">", "&gt;")
	return s
}

// SanitizeHTML normalizes tag nesting to prevent invalid HTML in Telegram parse mode.
func SanitizeHTML(input string) string {
	return normalizeHTMLTags(input)
}