3 files changed, 55 insertions, 2 deletions

diff --git a/resources/qml/MessageInput.qml b/resources/qml/MessageInput.qml
index c855ef90..9bb01471 100644
--- a/resources/qml/MessageInput.qml
+++ b/resources/qml/MessageInput.qml
@@ -209,6 +209,39 @@ Rectangle {
                     } else if (event.key == Qt.Key_Down && popup.opened) {
                         event.accepted = true;
                         popup.down();
+                    } else if (event.key == Qt.Key_Up) {
+                        if (cursorPosition == 0) {
+                            event.accepted = true;
+                            var idx = TimelineManager.timeline.edit ? TimelineManager.timeline.idToIndex(TimelineManager.timeline.edit) + 1 : 0;
+                            while (true) {
+                                var id = TimelineManager.timeline.indexToId(idx);
+                                if (!id || TimelineManager.timeline.getDump(id, "").isEditable) {
+                                    TimelineManager.timeline.edit = id;
+                                    cursorPosition = 0;
+                                    break;
+                                }
+                                idx++;
+                            }
+                        } else if (cursorPosition == messageInput.length) {
+                            event.accepted = true;
+                            cursorPosition = 0;
+                        }
+                    } else if (event.key == Qt.Key_Down) {
+                        if (cursorPosition == 0) {
+                            event.accepted = true;
+                            cursorPosition = messageInput.length;
+                        } else if (cursorPosition == messageInput.length && TimelineManager.timeline.edit) {
+                            event.accepted = true;
+                            var idx = TimelineManager.timeline.idToIndex(TimelineManager.timeline.edit) - 1;
+                            while (true) {
+                                var id = TimelineManager.timeline.indexToId(idx);
+                                if (!id || TimelineManager.timeline.getDump(id, "").isEditable) {
+                                    TimelineManager.timeline.edit = id;
+                                    break;
+                                }
+                                idx--;
+                            }
+                        }
                     }
                 }
                 background: null
diff --git a/src/timeline/TimelineModel.cpp b/src/timeline/TimelineModel.cpp
index af4c6aa2..d46a313a 100644
--- a/src/timeline/TimelineModel.cpp
+++ b/src/timeline/TimelineModel.cpp
@@ -499,6 +499,8 @@ TimelineModel::data(const mtx::events::collections::TimelineEvents &event, int r
                          data(event, static_cast<int>(ProportionalHeight)));
                 m.insert(names[Id], data(event, static_cast<int>(Id)));
                 m.insert(names[State], data(event, static_cast<int>(State)));
+                m.insert(names[IsEdited], data(event, static_cast<int>(IsEdited)));
+                m.insert(names[IsEditable], data(event, static_cast<int>(IsEditable)));
                 m.insert(names[IsEncrypted], data(event, static_cast<int>(IsEncrypted)));
                 m.insert(names[IsRoomEncrypted], data(event, static_cast<int>(IsRoomEncrypted)));
                 m.insert(names[ReplyTo], data(event, static_cast<int>(ReplyTo)));
@@ -1550,6 +1552,17 @@ TimelineModel::setEdit(QString newEdit)
         if (edit_.startsWith('m'))
                 return;
 
+        if (newEdit.isEmpty()) {
+                resetEdit();
+                return;
+        }
+
+        if (edit_.isEmpty()) {
+                this->textBeforeEdit  = input()->text();
+                this->replyBeforeEdit = reply_;
+                nhlog::ui()->debug("Stored: {}", textBeforeEdit.toStdString());
+        }
+
         if (edit_ != newEdit) {
                 auto ev = events.get(newEdit.toStdString(), "");
                 if (ev && mtx::accessors::sender(*ev) == http::client()->user_id().to_string()) {
@@ -1584,8 +1597,14 @@ TimelineModel::resetEdit()
         if (!edit_.isEmpty()) {
                 edit_ = "";
                 emit editChanged(edit_);
-                input()->setText("");
-                resetReply();
+                nhlog::ui()->debug("Restoring: {}", textBeforeEdit.toStdString());
+                input()->setText(textBeforeEdit);
+                textBeforeEdit.clear();
+                if (replyBeforeEdit.isEmpty())
+                        resetReply();
+                else
+                        setReply(replyBeforeEdit);
+                replyBeforeEdit.clear();
         }
 }
 
diff --git a/src/timeline/TimelineModel.h b/src/timeline/TimelineModel.h
index 5f599741..c4df300f 100644
--- a/src/timeline/TimelineModel.h
+++ b/src/timeline/TimelineModel.h
@@ -337,6 +337,7 @@ private:
 
         QString currentId, currentReadId;
         QString reply_, edit_;
+        QString textBeforeEdit, replyBeforeEdit;
         std::vector<QString> typingUsers_;
 
         TimelineViewManager *manager_;