From 4fb3886aab5eee59a4d9fec0b3caa25cd773edf8 Mon Sep 17 00:00:00 2001
From: Bohdan Horbeshko <bodqhrohro@gmail.com>
Date: Fri, 15 Aug 2025 04:31:27 -0400
Subject: [PATCH] Fix some MAM bugs

---
 telegram/utils.go | 11 +++++++----
 xmpp/handlers.go  | 48 +++++++++++++++++------------------------------
 2 files changed, 24 insertions(+), 35 deletions(-)

diff --git a/telegram/utils.go b/telegram/utils.go
index 004e9ef..8ea29a8 100644
--- a/telegram/utils.go
+++ b/telegram/utils.go
@@ -2387,11 +2387,11 @@ func (c *Client) GetMessagesBetween(chatID, fromMessageId, lastMessageId int64,
 	if limit < 0 {
 		reqFromMessageId = lastMessageId
 		reqOffset = -1
-		reqLimit = limit
+		reqLimit = -limit + 1
 	} else {
 		reqFromMessageId = fromMessageId
-		reqOffset = -limit
-		reqLimit = limit
+		reqOffset = -limit - 1
+		reqLimit = limit + 1
 	}
 
 	newMessages, err = c.client.GetChatHistory(&client.GetChatHistoryRequest{
@@ -2408,7 +2408,7 @@ func (c *Client) GetMessagesBetween(chatID, fromMessageId, lastMessageId int64,
 			fromPos := -1
 			for i, message := range newMessages.Messages {
 				if message.Id == fromMessageId {
-					fromPos = i
+					fromPos = i+1
 					break
 				}
 			}
@@ -2420,6 +2420,9 @@ func (c *Client) GetMessagesBetween(chatID, fromMessageId, lastMessageId int64,
 			complete = true
 		} else {
 			for _, message := range newMessages.Messages {
+				if message.Id == fromMessageId {
+					continue
+				}
 				if lastMessageId != 0 {
 					if message.Id > lastMessageId {
 						break
diff --git a/xmpp/handlers.go b/xmpp/handlers.go
index 9809385..0d0a6b4 100644
--- a/xmpp/handlers.go
+++ b/xmpp/handlers.go
@@ -2368,6 +2368,15 @@ func handleSetQueryMAM2(s xmpp.Sender, iq *stanza.IQ, query *extensions.MAM2Quer
 				iqAnswerSetError(answer, 404)
 				return
 			}
+			overallyLastMessage, overallyLastMessageErr := session.GetPreviousMessage(toID, beforeId)
+			if overallyLastMessageErr == nil && overallyLastMessage != nil {
+				overallyLastMessageId = overallyLastMessage.Id
+			}
+		} else {
+			newestMessage, newestMessageErr := session.GetPreviousMessage(toID, 0)
+			if newestMessageErr == nil && newestMessage != nil {
+				overallyLastMessageId = newestMessage.Id
+			}
 		}
 
 		if rsmAfter != 0 {
@@ -2392,14 +2401,13 @@ func handleSetQueryMAM2(s xmpp.Sender, iq *stanza.IQ, query *extensions.MAM2Quer
 			}
 			if rsmBefore < beforeId {
 				canBeComplete = false
-				overallyLastMessageId = beforeId
 			}
 			beforeId = rsmBefore
 		}
 
 		var lastMessageId int64
-		// should be fine even with afterId=0 as it would mean the last as needed
-		lastMessage, lastMessageErr := session.GetPreviousMessage(toID, afterId)
+		// should be fine even with beforeId=0 as it would mean the last as needed
+		lastMessage, lastMessageErr := session.GetPreviousMessage(toID, beforeId)
 		if lastMessageErr == nil && lastMessage != nil {
 			lastMessageId = lastMessage.Id
 		} else {
@@ -2411,16 +2419,6 @@ func handleSetQueryMAM2(s xmpp.Sender, iq *stanza.IQ, query *extensions.MAM2Quer
 			messages, newComplete, err = session.GetMessagesBetween(toID, afterId, lastMessageId, rsmLimit)
 			if canBeComplete {
 				complete = newComplete
-				if !complete {
-					if lastMessageId != 0 {
-						overallyLastMessageId = lastMessageId
-					} else {
-						newestMessage, newestMessageErr := session.GetPreviousMessage(toID, 0)
-						if newestMessageErr == nil && newestMessage != nil {
-							overallyLastMessageId = newestMessage.Id
-						}
-					}
-				}
 			}
 		}
 	} else { // time limit or no limits at all
@@ -2447,9 +2445,15 @@ func handleSetQueryMAM2(s xmpp.Sender, iq *stanza.IQ, query *extensions.MAM2Quer
 				endMsg, endMsgErr := session.FindMessageByTime(toID, endTime)
 				if endMsgErr == nil && endMsg != nil {
 					lastMessageId = endMsg.Id
+					overallyLastMessageId = lastMessageId
 				} else {
 					beyond = true
 				}
+			} else {
+				newestMessage, newestMessageErr := session.GetPreviousMessage(toID, 0)
+				if newestMessageErr == nil && newestMessage != nil {
+					overallyLastMessageId = newestMessage.Id
+				}
 			}
 
 			if rsmAfter != 0 {
@@ -2486,14 +2490,6 @@ func handleSetQueryMAM2(s xmpp.Sender, iq *stanza.IQ, query *extensions.MAM2Quer
 					if newLastMessageErr == nil && newLastMessage != nil {
 						if lastMessageId == 0 || lastMessageId != newLastMessage.Id {
 							canBeComplete = false
-							if lastMessageId != 0 {
-								overallyLastMessageId = lastMessageId
-							} else {
-								newestMessage, newestMessageErr := session.GetPreviousMessage(toID, 0)
-								if newestMessageErr == nil && newestMessage != nil {
-									overallyLastMessageId = newestMessage.Id
-								}
-							}
 							lastMessageId = newLastMessage.Id
 						}
 					} else {
@@ -2511,16 +2507,6 @@ func handleSetQueryMAM2(s xmpp.Sender, iq *stanza.IQ, query *extensions.MAM2Quer
 				messages, newComplete, err = session.GetMessagesBetween(toID, fromMessageId, lastMessageId, rsmLimit)
 				if canBeComplete {
 					complete = newComplete
-					if !complete {
-						if lastMessageId != 0 {
-							overallyLastMessageId = lastMessageId
-						} else {
-							newestMessage, newestMessageErr := session.GetPreviousMessage(toID, 0)
-							if newestMessageErr == nil && newestMessage != nil {
-								overallyLastMessageId = newestMessage.Id
-							}
-						}
-					}
 				}
 			}
 		}