From 9ac3433ddcbfc515667c55f68e3e5b3f386209c3 Mon Sep 17 00:00:00 2001
From: Team Humaki <team@humaki.com>
Date: Thu, 17 Sep 2026 02:38:06 -0700
Subject: [PATCH] parser: keep trailing comments after last section

---
 file.go      | 35 ++++++++++++++++++++++++++++++++++-
 file_test.go | 28 ++++++++++++++++++++++++++++
 parser.go    |  6 ++++++
 section.go   | 13 +++++++++----
 4 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/file.go b/file.go
index 58beabc..1d24a25 100644
--- a/file.go
+++ b/file.go
@@ -332,6 +332,32 @@ func (f *File) Append(source interface{}, others ...interface{}) error {
 	return f.Reload()
 }
 
+func writeCommentLines(buf *bytes.Buffer, comment, indent string) error {
+	if len(comment) == 0 {
+		return nil
+	}
+	lines := strings.Split(comment, LineBreak)
+	for i := range lines {
+		if len(lines[i]) == 0 {
+			continue
+		}
+		if len(indent) > 0 {
+			if _, err := buf.WriteString(indent); err != nil {
+				return err
+			}
+		}
+		if lines[i][0] != '#' && lines[i][0] != ';' {
+			lines[i] = "; " + strings.TrimSpace(lines[i])
+		} else {
+			lines[i] = lines[i][:1] + " " + strings.TrimSpace(lines[i][1:])
+		}
+		if _, err := buf.WriteString(lines[i] + LineBreak); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (f *File) writeToBuffer(indent string) (*bytes.Buffer, error) {
 	equalSign := DefaultFormatLeft + f.options.KeyValueDelimiterOnWrite + DefaultFormatRight
 
@@ -366,7 +392,7 @@ func (f *File) writeToBuffer(indent string) (*bytes.Buffer, error) {
 			}
 		} else {
 			// Write nothing if default section is empty
-			if len(sec.keyList) == 0 {
+			if len(sec.keyList) == 0 && len(sec.TrailingComment) == 0 {
 				continue
 			}
 		}
@@ -376,6 +402,9 @@ func (f *File) writeToBuffer(indent string) (*bytes.Buffer, error) {
 			if _, err := buf.WriteString(sec.rawBody); err != nil {
 				return nil, err
 			}
+			if err := writeCommentLines(buf, sec.TrailingComment, ""); err != nil {
+				return nil, err
+			}
 
 			if PrettySection && !isLastSection {
 				// Put a line between sections
@@ -495,6 +524,10 @@ func (f *File) writeToBuffer(indent string) (*bytes.Buffer, error) {
 			}
 		}
 
+		if err := writeCommentLines(buf, sec.TrailingComment, ""); err != nil {
+			return nil, err
+		}
+
 		if PrettySection && !isLastSection {
 			// Put a line between sections
 			if _, err := buf.WriteString(LineBreak); err != nil {
diff --git a/file_test.go b/file_test.go
index 306c0c8..e529ae8 100644
--- a/file_test.go
+++ b/file_test.go
@@ -19,6 +19,7 @@ import (
 	"os"
 	"runtime"
 	"sort"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -474,6 +475,33 @@ bar2 = "  val ue2 "
 bar3 = "  val ue3 "
 `, buf.String())
 	})
+
+	t.Run("keep trailing comments after last section on rewrite", func(t *testing.T) {
+		// Ansible blockinfile markers: # BEGIN before the section, # END after its keys.
+		src := `# BEGIN
+[profile test]
+sso_region = us-east-1
+output     = json
+# END
+`
+		f, err := Load([]byte(src))
+		require.NoError(t, err)
+		assert.Equal(t, "# BEGIN", f.Section("profile test").Comment)
+		assert.Equal(t, "# END", f.Section("profile test").TrailingComment)
+
+		_, err = f.NewSection("profile extra")
+		require.NoError(t, err)
+		f.Section("profile extra").Key("region").SetValue("us-west-2")
+
+		var buf bytes.Buffer
+		_, err = f.WriteTo(&buf)
+		require.NoError(t, err)
+		out := buf.String()
+		endIdx := strings.Index(out, "# END")
+		newIdx := strings.Index(out, "[profile extra]")
+		assert.True(t, endIdx >= 0 && newIdx > endIdx, "trailing # END should stay with the original section, got:\n%s", out)
+		assert.Contains(t, out, "[profile extra]")
+	})
 }
 
 func TestFile_SaveTo(t *testing.T) {
diff --git a/parser.go b/parser.go
index 513b19e..495e019 100644
--- a/parser.go
+++ b/parser.go
@@ -524,5 +524,11 @@ func (f *File) parse(reader io.Reader) (err error) {
 		p.comment.Reset()
 		lastRegularKey = key
 	}
+	// Comments after the last key/section are otherwise dropped on EOF
+	// (e.g. Ansible blockinfile # END markers). Keep them on the last section.
+	if p.comment.Len() > 0 {
+		section.TrailingComment = strings.TrimSpace(p.comment.String())
+		p.comment.Reset()
+	}
 	return nil
 }
diff --git a/section.go b/section.go
index a3615d8..c0f0d4a 100644
--- a/section.go
+++ b/section.go
@@ -24,10 +24,15 @@ import (
 type Section struct {
 	f        *File
 	Comment  string
-	name     string
-	keys     map[string]*Key
-	keyList  []string
-	keysHash map[string]string
+	// TrailingComment is comments after the last key of this section
+	// (or after a raw section body) through EOF. Comments that sit
+	// immediately before the next [section] still attach to that next
+	// section's Comment, matching existing parse behavior.
+	TrailingComment string
+	name            string
+	keys            map[string]*Key
+	keyList         []string
+	keysHash        map[string]string
 
 	isRawSection bool
 	rawBody      string