From c940c7b85be45d0b55568a9198c484608388ea95 Mon Sep 17 00:00:00 2001
From: sawka <mike.sawka@gmail.com>
Date: Tue, 11 Oct 2022 23:11:43 -0700
Subject: [PATCH] cmd-fg, update when cmd done back to input

---
 pkg/cmdrunner/cmdrunner.go |  2 +-
 pkg/remote/remote.go       |  6 +++++
 pkg/sstore/dbops.go        | 46 +++++++++++++++++++++++++++++++-------
 pkg/sstore/sstore.go       |  6 +++++
 4 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/pkg/cmdrunner/cmdrunner.go b/pkg/cmdrunner/cmdrunner.go
index f7999b17..733eb6f3 100644
--- a/pkg/cmdrunner/cmdrunner.go
+++ b/pkg/cmdrunner/cmdrunner.go
@@ -260,7 +260,7 @@ func RunCommand(ctx context.Context, pk *scpacket.FeCommandPacketType) (sstore.U
 	if sw != nil {
 		updateMap := make(map[string]interface{})
 		updateMap[sstore.SWField_SelectedLine] = rtnLine.LineNum
-		updateMap[sstore.SWField_Focus] = sstore.SWFocusCmd
+		updateMap[sstore.SWField_Focus] = sstore.SWFocusCmdFg
 		sw, err = sstore.UpdateScreenWindow(ctx, ids.SessionId, ids.ScreenId, ids.WindowId, updateMap)
 		if err != nil {
 			// ignore error again (nothing to do)
diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go
index ba7721e9..b1ea6983 100644
--- a/pkg/remote/remote.go
+++ b/pkg/remote/remote.go
@@ -1191,6 +1191,12 @@ func (msh *MShellProc) handleCmdDonePacket(donePk *packet.CmdDonePacketType) {
 		msh.WriteToPtyBuffer("[error] updating cmddone: %v\n", err)
 		return
 	}
+	sws, err := sstore.UpdateSWsWithCmdFg(context.Background(), donePk.CK.GetSessionId(), donePk.CK.GetCmdId())
+	if err != nil {
+		fmt.Printf("[error] trying to update cmd-fg screen windows: %v\n", err)
+		// fall-through (nothing to do)
+	}
+	update.ScreenWindows = sws
 	if update != nil {
 		sstore.MainBus.SendUpdate(donePk.CK.GetSessionId(), update)
 	}
diff --git a/pkg/sstore/dbops.go b/pkg/sstore/dbops.go
index 745d8665..bf3fbb55 100644
--- a/pkg/sstore/dbops.go
+++ b/pkg/sstore/dbops.go
@@ -655,7 +655,7 @@ func GetCmdById(ctx context.Context, sessionId string, cmdId string) (*CmdType,
 	return cmd, nil
 }
 
-func UpdateCmdDonePk(ctx context.Context, donePk *packet.CmdDonePacketType) (UpdatePacket, error) {
+func UpdateCmdDonePk(ctx context.Context, donePk *packet.CmdDonePacketType) (*ModelUpdate, error) {
 	if donePk == nil || donePk.CK.IsEmpty() {
 		return nil, fmt.Errorf("invalid cmddone packet (no ck)")
 	}
@@ -676,7 +676,7 @@ func UpdateCmdDonePk(ctx context.Context, donePk *packet.CmdDonePacketType) (Upd
 	if rtnCmd == nil {
 		return nil, fmt.Errorf("cmd data not found for ck[%s]", donePk.CK)
 	}
-	return ModelUpdate{Cmd: rtnCmd}, nil
+	return &ModelUpdate{Cmd: rtnCmd}, nil
 }
 
 func AppendCmdErrorPk(ctx context.Context, errPk *packet.CmdErrorPacketType) error {
@@ -690,12 +690,6 @@ func AppendCmdErrorPk(ctx context.Context, errPk *packet.CmdErrorPacketType) err
 	})
 }
 
-type SWKey struct {
-	SessionId string
-	ScreenId  string
-	WindowId  string
-}
-
 func HangupAllRunningCmds(ctx context.Context) error {
 	return WithTx(ctx, func(tx *TxWrap) error {
 		query := `UPDATE cmd SET status = ? WHERE status = ?`
@@ -1179,3 +1173,39 @@ func GetLineResolveItems(ctx context.Context, sessionId string, windowId string)
 	}
 	return rtn, nil
 }
+
+func UpdateSWsWithCmdFg(ctx context.Context, sessionId string, cmdId string) ([]*ScreenWindowType, error) {
+	var rtn []*ScreenWindowType
+	txErr := WithTx(ctx, func(tx *TxWrap) error {
+		query := `SELECT sessionid, screenid, windowid 
+                  FROM screen_window sw 
+                  WHERE 
+                    sessionid = ?
+                    AND focustype = 'cmd-fg' 
+                    AND selectedline IN (SELECT linenum 
+                                         FROM line l 
+                                         WHERE l.sessionid = sw.sessionid 
+                                           AND l.windowid = sw.windowid 
+                                           AND l.cmdid = ?
+                                        )`
+		var swKeys []SWKey
+		tx.SelectWrap(&swKeys, query, sessionId, cmdId)
+		if len(swKeys) == 0 {
+			return nil
+		}
+		for _, key := range swKeys {
+			editMap := make(map[string]interface{})
+			editMap[SWField_Focus] = SWFocusInput
+			sw, err := UpdateScreenWindow(tx.Context(), key.SessionId, key.ScreenId, key.WindowId, editMap)
+			if err != nil {
+				return err
+			}
+			rtn = append(rtn, sw)
+		}
+		return nil
+	})
+	if txErr != nil {
+		return nil, txErr
+	}
+	return rtn, nil
+}
diff --git a/pkg/sstore/sstore.go b/pkg/sstore/sstore.go
index 6d307544..2eb38a3f 100644
--- a/pkg/sstore/sstore.go
+++ b/pkg/sstore/sstore.go
@@ -394,6 +394,12 @@ func (a SWAnchorType) Value() (driver.Value, error) {
 	return quickValueJson(a)
 }
 
+type SWKey struct {
+	SessionId string
+	ScreenId  string
+	WindowId  string
+}
+
 type ScreenWindowType struct {
 	SessionId    string       `json:"sessionid"`
 	ScreenId     string       `json:"screenid"`