Transform mentions into MUC nicknames in incoming messages

telegram/formatter/formatter.go

@@ -2,6 +2,7 @@ package formatter
 
 import (
 	"sort"
+	"strings"
 	"unicode"
 
 	log "github.com/sirupsen/logrus"
@@ -25,9 +26,10 @@ const (
 
 // insertion is a piece of text in given position
 type insertion struct {
-	Offset int32
+	Offset    int32
-	Runes  []rune
+	Runes     []rune
-	Type   insertionType
+	Type      insertionType
+	Replacing bool
 }
 
 // insertionStack contains the sequence of insertions
@@ -74,6 +76,14 @@ func (s insertionStack) NewIterator() func() *insertion {
 	}
 }
 
+func isReplacing(entity *client.TextEntity) bool {
+	switch entity.Type.TextEntityTypeType() {
+	case client.TypeTextEntityTypeMention, client.TypeTextEntityTypeMentionName:
+		return true
+	}
+	return false
+}
+
 // SortEntities arranges the entities in traversal-ready order
 func SortEntities(entities []*client.TextEntity) []*client.TextEntity {
 	sortedEntities := make([]*client.TextEntity, len(entities))
@@ -85,6 +95,9 @@ func SortEntities(entities []*client.TextEntity) []*client.TextEntity {
 		if entity1.Offset < entity2.Offset {
 			return true
 		} else if entity1.Offset == entity2.Offset {
+			if entity1.Length == entity2.Length {
+				return !isReplacing(entity1) && isReplacing(entity2)
+			}
 			return entity1.Length > entity2.Length
 		}
 		return false
@@ -177,6 +190,22 @@ func ClaspDirectives(doubledRunes []rune, entities []*client.TextEntity) []*clie
 	return alignedEntities
 }
 
+func mentionBraces(entity *client.TextEntity, nickname string) []*insertion {
+	return []*insertion{
+		&insertion{
+			Offset:    entity.Offset,
+			Runes:     []rune(nickname),
+			Type:      insertionOpening,
+			Replacing: true,
+		},
+		&insertion{
+			Offset:    entity.Offset + entity.Length,
+			Type:      insertionClosing,
+			Replacing: true,
+		},
+	}
+}
+
 func markupBraces(entity *client.TextEntity, lbrace, rbrace []rune) []*insertion {
 	return []*insertion{
 		&insertion{
@@ -330,11 +359,37 @@ func textToDoubledRunes(text string) []rune {
 	return doubledRunes
 }
 
+// cuts a substring back from doubled runes
+func cutTextFromDoubledRunes(doubledRunes []rune, offset, length int32) string {
+	runeSlice := doubledRunes[offset:offset+length]
+	var str strings.Builder
+	var skipNext bool
+	for _, cp := range runeSlice {
+		if skipNext {
+			skipNext = false
+			continue
+		}
+
+		str.WriteRune(cp)
+
+		if cp > bmpCeil {
+			skipNext = true
+		}
+	}
+	return str.String()
+}
+
+type MentionRetriever interface {
+	GetMUCNicknameByUsername(username string) (string, error)
+	GetMUCNickname(id int64) string
+}
+
 // Format traverses an already sorted list of entities and wraps the text in a markup
 func Format(
 	sourceText string,
 	entities []*client.TextEntity,
 	markupMode MarkupModeType,
+	mentionRetriever MentionRetriever,
 ) string {
 	if len(entities) == 0 {
 		return sourceText
@@ -369,7 +424,23 @@ func Format(
 
 		startStack, endStack = startStack.rebalance(endStack, entity.Offset)
 
-		insertions := entityToMarkup(entity, doubledRunes, markupMode)
+		var insertions []*insertion
+		if entity != nil && entity.Type != nil {
+			switch entity.Type.TextEntityTypeType() {
+			case client.TypeTextEntityTypeMention:
+				username := cutTextFromDoubledRunes(doubledRunes, entity.Offset, entity.Length)
+				nickname, err := mentionRetriever.GetMUCNicknameByUsername(username)
+				if err == nil {
+					insertions = mentionBraces(entity, nickname)
+				}
+			case client.TypeTextEntityTypeMentionName:
+				mentionName, _ := entity.Type.(*client.TextEntityTypeMentionName)
+				nickname := mentionRetriever.GetMUCNickname(mentionName.UserId)
+				insertions = mentionBraces(entity, nickname)
+			default:
+				insertions = entityToMarkup(entity, doubledRunes, markupMode)
+			}
+		}
 		if len(insertions) > 1 {
 			startStack = append(startStack, insertions[0:len(insertions)-1]...)
 		}
@@ -417,6 +488,7 @@ func Format(
 	nextInsertion := startStack.NewIterator()
 	insertion := nextInsertion()
 	var skipNext bool
+	var insideReplacingEntity bool
 
 	for i, cp := range doubledRunes {
 		if skipNext {
@@ -424,19 +496,43 @@ func Format(
 			continue
 		}
 
+		// loop through possible multiple insertions at this point
 		for insertion != nil && int(insertion.Offset) <= i {
-			markupRunes = append(markupRunes, insertion.Runes...)
+			if !insideReplacingEntity {
+				markupRunes = append(markupRunes, insertion.Runes...)
+			}
+
+			// if replacing entity encountered, ignore all entities inside it until it's closed
+			// (replacing entities are assumed to be not nested or overlapped)
+			if insertion.Replacing {
+				if insertion.Type == insertionOpening {
+					insideReplacingEntity = true
+				} else if insertion.Type == insertionClosing {
+					insideReplacingEntity = false
+				}
+			}
+
 			insertion = nextInsertion()
 		}
 
+		if insideReplacingEntity {
+			continue
+		}
+
 		markupRunes = append(markupRunes, cp)
 		// skip two UTF-16 code units (not points actually!) if needed
 		if cp > bmpCeil {
 			skipNext = true
 		}
 	}
+	// flush closing insertions
 	for insertion != nil {
-		markupRunes = append(markupRunes, insertion.Runes...)
+		if !insideReplacingEntity {
+			markupRunes = append(markupRunes, insertion.Runes...)
+		}
+		if insertion.Replacing && insertion.Type == insertionClosing {
+			insideReplacingEntity = false
+		}
 		insertion = nextInsertion()
 	}
 

telegram/handlers.go

@@ -340,6 +340,7 @@ func (c *Client) updateMessageContent(update *client.UpdateMessageContent) {
 			textContent.Text.Text,
 			textContent.Text.Entities,
 			markupFunction,
+			c,
 		))
 
 		var from string

telegram/utils.go

@@ -1055,6 +1055,7 @@ func (c *Client) getMessageReply(message *client.Message, preview bool, noConten
 				replyTo.Quote.Text.Text,
 				replyTo.Quote.Text.Entities,
 				c.getFormatter(),
+				c,
 			)
 			// make the whole quote fit one line
 			text = strings.ReplaceAll(text, "\n", " ")
@@ -1392,6 +1393,7 @@ func (c *Client) messageContentToText(content client.MessageContent, chatId int6
 				photo.Caption.Text,
 				photo.Caption.Entities,
 				markupMode,
+				c,
 			)
 		}
 	case client.TypeMessageAudio:
@@ -1403,6 +1405,7 @@ func (c *Client) messageContentToText(content client.MessageContent, chatId int6
 				audio.Caption.Text,
 				audio.Caption.Entities,
 				markupMode,
+				c,
 			)
 		}
 	case client.TypeMessageVideo:
@@ -1414,6 +1417,7 @@ func (c *Client) messageContentToText(content client.MessageContent, chatId int6
 				video.Caption.Text,
 				video.Caption.Entities,
 				markupMode,
+				c,
 			)
 		}
 	case client.TypeMessageDocument:
@@ -1425,6 +1429,7 @@ func (c *Client) messageContentToText(content client.MessageContent, chatId int6
 				document.Caption.Text,
 				document.Caption.Entities,
 				markupMode,
+				c,
 			)
 		}
 	case client.TypeMessageText:
@@ -1436,6 +1441,7 @@ func (c *Client) messageContentToText(content client.MessageContent, chatId int6
 				text.Text.Text,
 				text.Text.Entities,
 				markupMode,
+				c,
 			)
 		}
 	case client.TypeMessageVoiceNote:
@@ -1447,6 +1453,7 @@ func (c *Client) messageContentToText(content client.MessageContent, chatId int6
 				voice.Caption.Text,
 				voice.Caption.Entities,
 				markupMode,
+				c,
 			)
 		}
 	case client.TypeMessageVideoNote:
@@ -1460,6 +1467,7 @@ func (c *Client) messageContentToText(content client.MessageContent, chatId int6
 				animation.Caption.Text,
 				animation.Caption.Entities,
 				markupMode,
+				c,
 			)
 		}
 	case client.TypeMessageContact:
@@ -2313,6 +2321,7 @@ func (c *Client) GetChatDescription(chat *client.Chat) string {
 					fullInfo.Bio.Text,
 					fullInfo.Bio.Entities,
 					c.getFormatter(),
+					c,
 				)
 			} else if fullInfo.BotInfo != nil {
 				if fullInfo.BotInfo.ShortDescription != "" {
@@ -3081,6 +3090,18 @@ func (c *Client) getChatMemberStatus(status client.ChatMemberStatus) ChatMemberS
 	return ChatMemberStatusIllegal
 }
 
+// GetMUCNicknameByUsername implement the MentionRetriever interface for message formatters
+func (c *Client) GetMUCNicknameByUsername(username string) (string, error) {
+	chat, err := c.client.SearchPublicChat(&client.SearchPublicChatRequest{
+		Username: username,
+	})
+	if err != nil {
+		return "", err
+	}
+
+	return c.GetMUCNickname(chat.Id), nil
+}
+
 // GetErrorCode obtains an error code from a Telegram response error
 func GetErrorCode(err error) (int32, bool) {
 	responseError, ok := err.(client.ResponseError)