telegram-mcp-read-only/cmd/test/unread.go

package main

import (
	"context"
	"sort"
	"time"

	"github.com/gotd/td/telegram"
	"github.com/gotd/td/tg"
	"github.com/pkg/errors"
	"github.com/rs/zerolog/log"
	cfg "github.com/spf13/pflag"
	"golang.org/x/time/rate"
)

const (
	defaultMessageLimit = 10
	maxDialogsLimit     = 100
	rateLimitPerSec     = 5
)

//nolint:gochecknoglobals // CLI flags must be global
var (
	messageLimit = cfg.Int("limit", defaultMessageLimit, "limit of unread messages to fetch")
)

//nolint:gochecknoglobals // Rate limiter should be global for consistent rate limiting across all functions
var telegramLimiter = rate.NewLimiter(rate.Limit(rateLimitPerSec), 1)

// UnreadMessage represents a simplified message structure
type UnreadMessage struct {
	ID        int
	Text      string
	Date      time.Time
	FromID    int64
	FromName  string
	ChatType  string
	ChatTitle string
}

// DialogWithUnread represents a dialog with its unread count and latest message ID
type DialogWithUnread struct {
	Dialog      *tg.Dialog
	UnreadCount int
	TopMessage  int
}

// getUnreadMessages fetches unread messages from different users
//
//nolint:gocognit,gocyclo // complexity is inherent to handling different types of Telegram messages and users
func getUnreadMessages(ctx context.Context, client *telegram.Client) ([]UnreadMessage, error) {
	if err := telegramLimiter.Wait(ctx); err != nil {
		return nil, errors.Wrap(err, "rate limiter wait")
	}

	api := client.API()
	dialogsClass, err := api.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{
		OffsetPeer:    &tg.InputPeerEmpty{},
		OffsetDate:    0,
		OffsetID:      0,
		Limit:         maxDialogsLimit,
		Hash:          0,
		Flags:         0,
		ExcludePinned: false,
		FolderID:      0,
	})
	if err != nil {
		return nil, errors.Wrap(err, "get dialogs")
	}

	var dialogs *tg.MessagesDialogs
	switch d := dialogsClass.(type) {
	case *tg.MessagesDialogs:
		dialogs = d
	case *tg.MessagesDialogsSlice:
		dialogs = &tg.MessagesDialogs{
			Dialogs:  d.Dialogs,
			Messages: d.Messages,
			Chats:    d.Chats,
			Users:    d.Users,
		}
	default:
		return nil, errors.New("unexpected dialogs response type")
	}

	// Create a slice of dialogs with unread count
	dialogsWithUnread := make([]DialogWithUnread, 0, len(dialogs.Dialogs))
	for _, dialog := range dialogs.Dialogs {
		dialogItem, ok := dialog.(*tg.Dialog)
		if !ok {
			continue
		}

		if dialogItem.UnreadCount > 0 {
			dialogsWithUnread = append(dialogsWithUnread, DialogWithUnread{
				Dialog:      dialogItem,
				UnreadCount: dialogItem.UnreadCount,
				TopMessage:  dialogItem.TopMessage,
			})
		}
	}

	// Sort dialogs by TopMessage in descending order (newest first)
	sort.Slice(dialogsWithUnread, func(i, j int) bool {
		return dialogsWithUnread[i].TopMessage > dialogsWithUnread[j].TopMessage
	})

	// Map to store the latest message from each user
	userMessages := make(map[int64]UnreadMessage)
	processedCount := 0

	for _, dialogWithUnread := range dialogsWithUnread {
		dialogItem := dialogWithUnread.Dialog

		var inputPeer tg.InputPeerClass
		var chatType, chatTitle string
		var fromID int64
		var fromName string

		switch peer := dialogItem.Peer.(type) {
		case *tg.PeerUser:
			for _, userItem := range dialogs.Users {
				user, ok := userItem.(*tg.User)
				if !ok || user.ID != peer.UserID {
					continue
				}

				inputPeer = &tg.InputPeerUser{
					UserID:     user.ID,
					AccessHash: user.AccessHash,
				}
				chatType = "user"
				chatTitle = user.FirstName + " " + user.LastName
				fromID = user.ID
				fromName = chatTitle

				break
			}
		case *tg.PeerChat:
			inputPeer = &tg.InputPeerChat{
				ChatID: peer.ChatID,
			}
			chatType = "chat"
			for _, chatItem := range dialogs.Chats {
				chat, ok := chatItem.(*tg.Chat)
				if !ok || chat.ID != peer.ChatID {
					continue
				}

				chatTitle = chat.Title

				break
			}
		case *tg.PeerChannel:
			for _, channelItem := range dialogs.Chats {
				channel, ok := channelItem.(*tg.Channel)
				if !ok || channel.ID != peer.ChannelID {
					continue
				}

				inputPeer = &tg.InputPeerChannel{
					ChannelID:  channel.ID,
					AccessHash: channel.AccessHash,
				}
				chatType = "channel"
				chatTitle = channel.Title

				break
			}
		}

		if inputPeer == nil {
			continue
		}

		if err := telegramLimiter.Wait(ctx); err != nil {
			return nil, errors.Wrap(err, "rate limiter wait")
		}

		messagesClass, err := api.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{
			Peer:       inputPeer,
			OffsetID:   0,
			OffsetDate: 0,
			AddOffset:  0,
			Limit:      1, // We only need the latest message
			MaxID:      0,
			MinID:      0,
			Hash:       0,
		})
		if err != nil {
			log.Error().Err(err).Msg("failed to get messages")

			continue
		}

		var messages *tg.MessagesMessages
		switch m := messagesClass.(type) {
		case *tg.MessagesMessages:
			messages = m
		case *tg.MessagesMessagesSlice:
			messages = &tg.MessagesMessages{
				Messages: m.Messages,
				Chats:    m.Chats,
				Users:    m.Users,
			}
		case *tg.MessagesChannelMessages:
			messages = &tg.MessagesMessages{
				Messages: m.Messages,
				Chats:    m.Chats,
				Users:    m.Users,
			}
		default:
			log.Error().Msg("unexpected messages response type")

			continue
		}

		for _, msg := range messages.Messages {
			message, ok := msg.(*tg.Message)
			if !ok {
				continue
			}

			if message.Out {
				continue
			}

			if message.FromID != nil {
				if from, ok := message.FromID.(*tg.PeerUser); ok {
					for _, userItem := range messages.Users {
						user, ok := userItem.(*tg.User)
						if !ok || user.ID != from.UserID {
							continue
						}

						fromID = user.ID
						fromName = user.FirstName + " " + user.LastName

						break
					}
				}
			}

			unreadMsg := UnreadMessage{
				ID:        message.ID,
				Text:      message.Message,
				Date:      time.Unix(int64(message.Date), 0),
				FromID:    fromID,
				FromName:  fromName,
				ChatType:  chatType,
				ChatTitle: chatTitle,
			}

			// Only store if we haven't seen this user yet or if this message is newer
			if existingMsg, exists := userMessages[fromID]; !exists || unreadMsg.Date.After(existingMsg.Date) {
				userMessages[fromID] = unreadMsg
				processedCount++
			}

			break // We only need the latest message
		}

		if len(userMessages) >= *messageLimit {
			break
		}
	}

	// Convert map to slice and sort by date
	messages := make([]UnreadMessage, 0, len(userMessages))
	for _, msg := range userMessages {
		messages = append(messages, msg)
	}

	// Sort messages by date in descending order
	sort.Slice(messages, func(i, j int) bool {
		return messages[i].Date.After(messages[j].Date)
	})

	return messages, nil
}