From dbd82e49395c8781f5f7fe754f77be44d1ee73ce Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sat, 21 Sep 2024 15:30:34 +0200
Subject: [PATCH] fix(common): Improve formatting

---
 common/formatting/formatting.go      | 41 ++++++++++++++++++++++++++++
 common/formatting/formatting_test.go | 37 +++++++++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/common/formatting/formatting.go b/common/formatting/formatting.go
index 3400aaf..c72b639 100644
--- a/common/formatting/formatting.go
+++ b/common/formatting/formatting.go
@@ -32,6 +32,8 @@ func (f FormatTemplate) Format(
 		value = strings.TrimRight(value, "\t")
 	}
 
+	value = surroundWithQuotes(value)
+
 	return value
 }
 
@@ -57,6 +59,45 @@ func (f FormatTemplate) replace(format string, replacement string) FormatTemplat
 	return FormatTemplate(value)
 }
 
+func surroundWithQuotes(s string) string {
+	value := s
+	currentIndex := 0
+
+	for {
+		startPosition := strings.Index(value[currentIndex:], "/!'")
+
+		if startPosition == -1 {
+			break
+		}
+
+		startPosition = startPosition + currentIndex + 3
+		currentIndex = startPosition
+
+		endPosition := strings.Index(value[startPosition:], "/!'")
+
+		if endPosition == -1 {
+			break
+		}
+
+		endPosition = endPosition + startPosition
+		currentIndex = endPosition
+
+		innerValue := value[startPosition:endPosition]
+
+		if strings.Contains(innerValue, " ") {
+			value = value[:startPosition-3] + "\"" + innerValue + "\"" + value[endPosition+3:]
+		} else {
+			value = value[:startPosition-3] + innerValue + value[endPosition+3:]
+		}
+
+		if endPosition+3 >= len(value) {
+			break
+		}
+	}
+
+	return value
+}
+
 func getTab(options protocol.FormattingOptions) string {
 	tabSize := options["tabSize"].(float64)
 	insertSpace := options["insertSpaces"].(bool)
diff --git a/common/formatting/formatting_test.go b/common/formatting/formatting_test.go
index 290d41f..fc606aa 100644
--- a/common/formatting/formatting_test.go
+++ b/common/formatting/formatting_test.go
@@ -79,3 +79,40 @@ func TestSimpleExampleWhiteSpaceAtEndShouldNOTTrim(
 		t.Errorf("Expected %q but got %q", expected, result)
 	}
 }
+
+func TestSurroundWithQuotesExample(
+	t *testing.T,
+) {
+	template := FormatTemplate("%s /!'%s/!'")
+
+	options := protocol.FormattingOptions{
+		"tabSize":                float64(4),
+		"insertSpaces":           false,
+		"trimTrailingWhitespace": true,
+	}
+
+	result := template.Format(options, "PermitRootLogin", "this is okay")
+	expected := `PermitRootLogin "this is okay"`
+
+	if result != expected {
+		t.Errorf("Expected %q but got %q", expected, result)
+	}
+}
+func TestSurroundWithQuotesButNoSpaceExample(
+	t *testing.T,
+) {
+	template := FormatTemplate("%s /!'%s/!'")
+
+	options := protocol.FormattingOptions{
+		"tabSize":                float64(4),
+		"insertSpaces":           false,
+		"trimTrailingWhitespace": true,
+	}
+
+	result := template.Format(options, "PermitRootLogin", "yes")
+	expected := `PermitRootLogin yes`
+
+	if result != expected {
+		t.Errorf("Expected %q but got %q", expected, result)
+	}
+}