2 files changed, 49 insertions, 51 deletions

diff --git a/worker/notmuch/message.go b/worker/notmuch/message.go
index 6529507..3d3ed0c 100644
--- a/worker/notmuch/message.go
+++ b/worker/notmuch/message.go
@@ -64,70 +64,68 @@ func (m *Message) NewBodyPartReader(requestedParts []int) (io.Reader, error) {
 	return lib.FetchEntityPartReader(msg, requestedParts)
 }
 
-// MarkAnswered either adds or removes the "replied" tag from the message.
-func (m *Message) MarkAnswered(answered bool) error {
-	haveReplied := false
+// SetOneFlag adds or removes a single flag from the message.
+// Notmuch doesn't support all the flags, and for those this errors.
+func (m *Message) SetOneFlag(flag models.Flag, enable bool) error {
+	// Translate the flag into a notmuch tag, ignoring no-op flags.
+	var tag string
+	switch flag {
+	case models.SeenFlag:
+		// Note: Inverted properly later
+		tag = "unread"
+	case models.AnsweredFlag:
+		tag = "replied"
+	case models.FlaggedFlag:
+		tag = "flagged"
+	default:
+		return fmt.Errorf("Notmuch doesn't support flag %v", flag)
+	}
+
+	// Get the current state of the flag.
+	// Note that notmuch handles models.SeenFlag in an inverted sense.
+	oldState := false
 	tags, err := m.Tags()
 	if err != nil {
 		return err
 	}
 	for _, t := range tags {
-		if t == "replied" {
-			haveReplied = true
+		if t == tag {
+			oldState = true
 			break
 		}
 	}
-	if haveReplied == answered {
-		// we already have the desired state
-		return nil
+	if flag == models.SeenFlag {
+		oldState = !oldState
 	}
 
-	if haveReplied {
-		err := m.RemoveTag("replied")
-		if err != nil {
-			return err
-		}
+	// Skip if flag already in correct state.
+	if oldState == enable {
 		return nil
 	}
 
-	err = m.AddTag("replied")
-	if err != nil {
-		return err
+	if !enable {
+		if flag == models.SeenFlag {
+			return m.AddTag("unread")
+		} else {
+			return m.RemoveTag(tag)
+		}
+	} else {
+		if flag == models.SeenFlag {
+			return m.RemoveTag("unread")
+		} else {
+			return m.AddTag(tag)
+		}
 	}
-	return nil
+}
+
+// MarkAnswered either adds or removes the "replied" tag from the message.
+func (m *Message) MarkAnswered(answered bool) error {
+	return m.SetOneFlag(models.AnsweredFlag, answered)
 }
 
 // MarkRead either adds or removes the maildir.FlagSeen flag from the message.
 func (m *Message) MarkRead(seen bool) error {
-	haveUnread := false
-	tags, err := m.Tags()
-	if err != nil {
-		return err
-	}
-	for _, t := range tags {
-		if t == "unread" {
-			haveUnread = true
-			break
-		}
-	}
-	if (haveUnread && !seen) || (!haveUnread && seen) {
-		// we already have the desired state
-		return nil
-	}
-
-	if haveUnread {
-		err := m.RemoveTag("unread")
-		if err != nil {
-			return err
-		}
-		return nil
-	}
-
-	err = m.AddTag("unread")
-	if err != nil {
-		return err
-	}
-	return nil
+	return m.SetOneFlag(models.SeenFlag, seen)
 }
 
 // tags returns the notmuch tags of a message
diff --git a/worker/notmuch/worker.go b/worker/notmuch/worker.go
index 3fdd11d..f14b7ff 100644
--- a/worker/notmuch/worker.go
+++ b/worker/notmuch/worker.go
@@ -113,8 +113,8 @@ func (w *worker) handleMessage(msg types.WorkerMessage) error {
 		return w.handleFetchMessageBodyPart(msg)
 	case *types.FetchFullMessages:
 		return w.handleFetchFullMessages(msg)
-	case *types.ReadMessages:
-		return w.handleReadMessages(msg)
+	case *types.FlagMessages:
+		return w.handleFlagMessages(msg)
 	case *types.AnsweredMessages:
 		return w.handleAnsweredMessages(msg)
 	case *types.SearchDirectory:
@@ -392,7 +392,7 @@ func (w *worker) handleAnsweredMessages(msg *types.AnsweredMessages) error {
 	return nil
 }
 
-func (w *worker) handleReadMessages(msg *types.ReadMessages) error {
+func (w *worker) handleFlagMessages(msg *types.FlagMessages) error {
 	for _, uid := range msg.Uids {
 		m, err := w.msgFromUid(uid)
 		if err != nil {
@@ -400,8 +400,8 @@ func (w *worker) handleReadMessages(msg *types.ReadMessages) error {
 			w.err(msg, err)
 			continue
 		}
-		if err := m.MarkRead(msg.Read); err != nil {
-			w.w.Logger.Printf("could not mark message as read: %v", err)
+		if err := m.SetOneFlag(msg.Flag, msg.Enable); err != nil {
+			w.w.Logger.Printf("could not set flag %v as %v for message: %v", msg.Flag, msg.Enable, err)
 			w.err(msg, err)
 			continue
 		}