From ec6858ff976ff8b90a9a1a609315a0fe926255e8 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Wed, 4 Sep 2024 15:21:05 +0100
Subject: [PATCH] request: refactor prepareRequests()

---
 event/event.go  |   2 +-
 request.go      | 118 ++++++++++++++++++++++++------------------------
 request_test.go |   4 +-
 3 files changed, 63 insertions(+), 61 deletions(-)

diff --git a/event/event.go b/event/event.go
index 26f6764..69e9f72 100644
--- a/event/event.go
+++ b/event/event.go
@@ -148,7 +148,7 @@ func processEvent(ev EventHandler, msg ReceivedData, events []EventType) {
 					WorkspaceName: WorkspaceName(raw[1]),
 				})
 			case EventActiveWindow:
-				// e.g. jetbrains-goland,hyprland-ipc-ipc – main.go
+				// e.g. nvim,nvim event/event.go
 				ev.ActiveWindow(ActiveWindow{
 					Name:  raw[0],
 					Title: raw[1],
diff --git a/request.go b/request.go
index 932ef88..f3b58a5 100644
--- a/request.go
+++ b/request.go
@@ -24,19 +24,41 @@ const (
 var reqHeader = []byte{'j', '/'}
 var reqSep = []byte{' ', ';'}
 
-func prepareRequest(buf *bytes.Buffer, command string, param string, jsonResp bool) int {
+func prepareRequest(
+	buf *bytes.Buffer,
+	command string,
+	param string,
+	jsonResp bool,
+	lastReq bool,
+) (int, error) {
+	prevLen := buf.Len()
+
 	if jsonResp {
 		buf.Write(reqHeader)
 	}
 	buf.WriteString(command)
-	buf.WriteByte(reqSep[0])
-	buf.WriteString(param)
-	buf.WriteByte(reqSep[1])
+	if param != "" {
+		buf.WriteByte(reqSep[0])
+		buf.WriteString(param)
+	}
+	if !lastReq {
+		buf.WriteByte(reqSep[1])
+	}
+
+	reqLen := buf.Len() - prevLen
+
+	if buf.Len() > bufSize {
+		return reqLen, fmt.Errorf(
+			"command is too long (%d>=%d): %s",
+			buf.Len(),
+			bufSize,
+			buf.String(),
+		)
+	}
 
-	return buf.Len()
+	return reqLen, nil
 }
 
-// TODO: needs refactor, the logic is all over the place
 func prepareRequests(command string, params []string, jsonResp bool) (requests []RawRequest, err error) {
 	if command == "" {
 		// Panic since this is not supposed to happen, i.e.: only by
@@ -48,67 +70,47 @@ func prepareRequests(command string, params []string, jsonResp bool) (requests [
 	buf := bytes.NewBuffer(nil)
 
 	switch len(params) {
-	case 0, 1:
-		if jsonResp {
-			buf.Write(reqHeader)
-		}
-		buf.WriteString(command)
-		if len(params) == 1 {
-			buf.WriteByte(reqSep[0])
-			buf.WriteString(params[0])
+	case 0:
+		_, err := prepareRequest(buf, command, "", jsonResp, true)
+		if err != nil {
+			return nil, err
 		}
-
-		if buf.Len() > bufSize {
-			return nil, fmt.Errorf(
-				"command is too long (%d>=%d): %s",
-				buf.Len(),
-				bufSize,
-				buf.String(),
-			)
+	case 1:
+		_, err := prepareRequest(buf, command, params[0], jsonResp, true)
+		if err != nil {
+			return nil, err
 		}
 	default:
 		// Add [[BATCH]] to the buffer
 		buf.WriteString(batch)
-		// Initialise current length of buffer
-		curLen := buf.Len()
-
-		for _, param := range params {
-			// Get the current command + param length + request
-			// header and separators
-			cmdLen := len(command) + len(param) + len(reqSep)
-			if jsonResp {
-				cmdLen += len(reqHeader)
-			}
-
-			// If batch + command length is bigger than bufSize,
-			// return an error since it will not fit the socket
-			if len(batch)+cmdLen > bufSize {
-				return nil, fmt.Errorf(
-					"command is too long (%d>=%d): %s%s %s;",
-					len(batch)+cmdLen,
-					bufSize,
-					batch,
-					command,
-					param,
-				)
-			}
 
-			// If the current length of the buffer + command +
-			// param is bigger than bufSize, we will need to split
-			// the request
-			if curLen+cmdLen > bufSize {
-				// Append current buffer contents to the
-				// requests array
-				requests = append(requests, buf.Bytes())
-
-				// Reset the current buffer and add [[BATCH]]
-				buf.Reset()
-				buf.WriteString(batch)
+		for i, param := range params {
+		prepare:
+			reqLen, err := prepareRequest(buf, command, param, jsonResp, i == len(params)-1)
+			if err != nil {
+				// If request length + [[BATCH]] < bufSize, the
+				// request will fit as long as we reset the
+				// buffer
+				if reqLen+len(batch) <= bufSize {
+					// Append current buffer contents to
+					// the requests array, ignoring the
+					// current request
+					buf.Truncate(bufSize - reqLen)
+					requests = append(requests, buf.Bytes())
+
+					// Reset the current buffer and add
+					// [[BATCH]]
+					buf.Reset()
+					buf.WriteString(batch)
+
+					// Prepare the current request again
+					goto prepare
+				}
+				return nil, err
 			}
-			// Add the contents of the request to the buffer
-			curLen = prepareRequest(buf, command, param, jsonResp)
 		}
 	}
+
 	// Append any remaining buffer content to requests array
 	requests = append(requests, buf.Bytes())
 
diff --git a/request_test.go b/request_test.go
index 288fdf5..00f315f 100644
--- a/request_test.go
+++ b/request_test.go
@@ -78,10 +78,10 @@ func TestPrepareRequests(t *testing.T) {
 	}{
 		{"command", nil, true, []string{"j/command"}},
 		{"command", []string{"param0"}, true, []string{"j/command param0"}},
-		{"command", []string{"param0", "param1"}, true, []string{"[[BATCH]]j/command param0;j/command param1;"}},
+		{"command", []string{"param0", "param1"}, true, []string{"[[BATCH]]j/command param0;j/command param1"}},
 		{"command", nil, false, []string{"command"}},
 		{"command", []string{"param0"}, false, []string{"command param0"}},
-		{"command", []string{"param0", "param1"}, false, []string{"[[BATCH]]command param0;command param1;"}},
+		{"command", []string{"param0", "param1"}, false, []string{"[[BATCH]]command param0;command param1"}},
 	}
 	for _, tt := range tests {
 		t.Run(fmt.Sprintf("tests_%s-%s-%v", tt.command, tt.params, tt.jsonResp), func(t *testing.T) {