2 files changed, 90 insertions, 18 deletions

diff --git a/matrix/rooms/room.go b/matrix/rooms/room.go
index 40303be..17bf21b 100644
--- a/matrix/rooms/room.go
+++ b/matrix/rooms/room.go
@@ -43,6 +43,12 @@ type RoomTag struct {
 	Order string
 }
 
+type UnreadMessage struct {
+	EventID string
+	Counted bool
+	Highlight bool
+}
+
 // Room represents a single Matrix room.
 type Room struct {
 	*gomatrix.Room
@@ -57,13 +63,11 @@ type Room struct {
 	SessionUserID string
 
 	// The number of unread messages that were notified about.
-	UnreadMessages int
-	// Whether or not any of the unread messages were highlights.
-	Highlighted bool
-	// Whether or not the room contains any new messages.
-	// This can be true even when UnreadMessages is zero if there's
-	// a notificationless message like bot notices.
-	HasNewMessages bool
+	UnreadMessages []UnreadMessage
+	unreadCountCache *int
+	highlightCache *bool
+	// Whether or not this room is marked as a direct chat.
+	IsDirect bool
 
 	// List of tags given to this room
 	RawTags []RoomTag
@@ -110,14 +114,74 @@ func (room *Room) UnlockHistory() {
 }
 
 // MarkRead clears the new message statuses on this room.
-func (room *Room) MarkRead() {
-	room.UnreadMessages = 0
-	room.Highlighted = false
-	room.HasNewMessages = false
+func (room *Room) MarkRead(eventID string) {
+	readToIndex := -1
+	for index, unreadMessage := range room.UnreadMessages {
+		if unreadMessage.EventID == eventID {
+			readToIndex = index
+		}
+	}
+	if readToIndex >= 0 {
+		room.UnreadMessages = room.UnreadMessages[readToIndex+1:]
+		room.highlightCache = nil
+		room.unreadCountCache = nil
+	}
+}
+
+func (room *Room) UnreadCount() int {
+	if room.unreadCountCache == nil {
+		room.unreadCountCache = new(int)
+		for _, unreadMessage := range room.UnreadMessages {
+			if unreadMessage.Counted {
+				*room.unreadCountCache++
+			}
+		}
+	}
+	return *room.unreadCountCache
+}
+
+func (room *Room) Highlighted() bool {
+	if room.highlightCache == nil {
+		room.highlightCache = new(bool)
+		for _, unreadMessage := range room.UnreadMessages {
+			if unreadMessage.Highlight {
+				*room.highlightCache = true
+				break
+			}
+		}
+	}
+	return *room.highlightCache
+}
+
+func (room *Room) HasNewMessages() bool {
+	return len(room.UnreadMessages) > 0
+}
+
+func (room *Room) AddUnread(eventID string, counted, highlight bool) {
+	room.UnreadMessages = append(room.UnreadMessages, UnreadMessage{
+		EventID: eventID,
+		Counted: counted,
+		Highlight: highlight,
+	})
+	if counted {
+		if room.unreadCountCache == nil {
+			room.unreadCountCache = new(int)
+		}
+		*room.unreadCountCache++
+	}
+	if highlight {
+		if room.highlightCache == nil {
+			room.highlightCache = new(bool)
+		}
+		*room.highlightCache = true
+	}
 }
 
 func (room *Room) Tags() []RoomTag {
 	if len(room.RawTags) == 0 {
+		if room.IsDirect {
+			return []RoomTag{{"net.maunium.gomuks.fake.direct", "0.5"}}
+		}
 		return []RoomTag{{"", "0.5"}}
 	}
 	return room.RawTags
diff --git a/matrix/rooms/room_test.go b/matrix/rooms/room_test.go
index 258e57b..1bd47ff 100644
--- a/matrix/rooms/room_test.go
+++ b/matrix/rooms/room_test.go
@@ -215,11 +215,19 @@ func TestRoom_GetTitle_Members_GroupChat(t *testing.T) {
 
 func TestRoom_MarkRead(t *testing.T) {
 	room := rooms.NewRoom("!test:maunium.net", "@tulir:maunium.net")
-	room.UnreadMessages = 123
-	room.Highlighted = true
-	room.HasNewMessages = true
-	room.MarkRead()
-	assert.Zero(t, room.UnreadMessages)
-	assert.False(t, room.Highlighted)
-	assert.False(t, room.HasNewMessages)
+
+	room.AddUnread("foo", true, false)
+	assert.Equal(t, 1, room.UnreadCount())
+	assert.False(t, room.Highlighted())
+
+	room.AddUnread("bar", true, false)
+	assert.Equal(t, 2, room.UnreadCount())
+	assert.False(t, room.Highlighted())
+
+	room.AddUnread("asd", false, true)
+	assert.Equal(t, 2, room.UnreadCount())
+	assert.True(t, room.Highlighted())
+
+	room.MarkRead("")
+	assert.Empty(t, room.UnreadMessages)
 }