From b71d24a0f0de5cd211eaf9852e55722952bc86ff Mon Sep 17 00:00:00 2001
From: Anirudh Sudhir <anirudh.sudhir1@gmail.com>
Date: Tue, 9 Apr 2024 00:06:08 +0530
Subject: [PATCH] feat: auto-refresh browser on live reload

TODO: If the browser is refreshed manually, the file has to be modified
  twice to reload the page later.

Co-authored-by: Aditya Hegde <aditya.mh@outlook.com>
---
 cmd/anna/anna.go               |  2 ++
 cmd/anna/livereload.go         | 35 +++++++++++++++++++++++++++++++++-
 main.go                        |  1 +
 pkg/parser/parser.go           |  6 +++++-
 site/layout/page.html          |  2 +-
 site/layout/partials/head.html | 18 ++++++++++++++++-
 6 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/cmd/anna/anna.go b/cmd/anna/anna.go
index dc943dd..be9ea6b 100644
--- a/cmd/anna/anna.go
+++ b/cmd/anna/anna.go
@@ -14,6 +14,7 @@ import (
 type Cmd struct {
 	RenderDrafts bool
 	Addr         string
+	LiveReload   bool
 }
 
 func (cmd *Cmd) VanillaRender() {
@@ -23,6 +24,7 @@ func (cmd *Cmd) VanillaRender() {
 		TagsMap:      make(map[string][]parser.TemplateData),
 		ErrorLogger:  log.New(os.Stderr, "ERROR\t", log.Ldate|log.Ltime|log.Lshortfile),
 		RenderDrafts: cmd.RenderDrafts,
+		LiveReload:   cmd.LiveReload,
 	}
 	e := engine.Engine{
 		Templates:   make(map[template.URL]parser.TemplateData),
diff --git a/cmd/anna/livereload.go b/cmd/anna/livereload.go
index ecf2b27..00a133f 100644
--- a/cmd/anna/livereload.go
+++ b/cmd/anna/livereload.go
@@ -6,11 +6,16 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"sync/atomic"
 	"time"
 
 	"github.com/acmpesuecc/anna/pkg/helpers"
 )
 
+var reloadPage = make(chan struct{})
+
+var countRequests atomic.Int32
+
 type liveReload struct {
 	errorLogger *log.Logger
 	fileTimes   map[string]time.Time
@@ -43,6 +48,7 @@ func (cmd *Cmd) StartLiveReload() {
 		for _, rootDir := range lr.rootDirs {
 			if lr.traverseDirectory(rootDir) {
 				cmd.VanillaRender()
+				reloadPage <- struct{}{}
 			}
 		}
 		if !lr.serverRunning {
@@ -97,8 +103,35 @@ func (lr *liveReload) checkFile(path string, modTime time.Time) bool {
 
 func (lr *liveReload) startServer(addr string) {
 	fmt.Print("Serving content at: http://localhost:", addr, "\n")
-	err := http.ListenAndServe(":"+addr, http.FileServer(http.Dir(helpers.SiteDataPath+"./rendered")))
+	http.Handle("/", http.FileServer(http.Dir(helpers.SiteDataPath+"./rendered")))
+	http.HandleFunc("/events", eventsHandler)
+	err := http.ListenAndServe(":"+addr, nil)
 	if err != nil {
 		lr.errorLogger.Fatal(err)
 	}
 }
+
+func eventsHandler(w http.ResponseWriter, r *http.Request) {
+	countRequests.Add(1)
+
+	// Set CORS headers to allow all origins.
+	w.Header().Set("Access-Control-Allow-Origin", "*")
+	w.Header().Set("Access-Control-Expose-Headers", "Content-Type")
+
+	w.Header().Set("Content-Type", "text/event-stream")
+	w.Header().Set("Cache-Control", "no-cache")
+	w.Header().Set("Connection", "keep-alive")
+
+	if countRequests.Load() == 1 {
+		<-reloadPage
+	} else {
+		countRequests.Store(countRequests.Load() - 1)
+		return
+	}
+
+	event := "event:\ndata:\n\n"
+	w.Write([]byte(event))
+	w.(http.Flusher).Flush()
+
+	countRequests.Store(countRequests.Load() - 1)
+}
diff --git a/main.go b/main.go
index 6ece361..d84b4f3 100644
--- a/main.go
+++ b/main.go
@@ -29,6 +29,7 @@ func main() {
 			}
 
 			if serve {
+				annaCmd.LiveReload = true
 				annaCmd.StartLiveReload()
 			}
 
diff --git a/pkg/parser/parser.go b/pkg/parser/parser.go
index f0d150b..9a4b6f2 100644
--- a/pkg/parser/parser.go
+++ b/pkg/parser/parser.go
@@ -46,6 +46,7 @@ type TemplateData struct {
 	Frontmatter              Frontmatter
 	Body                     template.HTML
 	Layout                   LayoutConfig
+	LiveReload               bool
 
 	// Do not use these fields to store tags!!
 	// These fields are populated by the ssg to store merged tag data
@@ -80,6 +81,9 @@ type Parser struct {
 	ErrorLogger *log.Logger
 
 	Helper *helpers.Helper
+
+	// Determines the injection of Live Reload JS in HTML
+	LiveReload bool
 }
 
 func (p *Parser) ParseMDDir(baseDirPath string, baseDirFS fs.FS) {
@@ -135,6 +139,7 @@ func (p *Parser) AddFileAndRender(baseDirPath string, dirEntryPath string, front
 		Frontmatter:              frontmatter,
 		Body:                     template.HTML(body),
 		Layout:                   p.LayoutConfig,
+		LiveReload:               p.LiveReload,
 	}
 
 	// Adding the page to the merged map storing all site pages
@@ -253,7 +258,6 @@ func (p *Parser) ParseRobots(inFilePath string, outFilePath string) {
 
 // Parse all the ".html" layout files in the layout/ directory
 func (p *Parser) ParseLayoutFiles() *template.Template {
-
 	// Parsing all files in the layout/ dir which match the "*.html" pattern
 	templ, err := template.ParseGlob(helpers.SiteDataPath + "layout/*.html")
 	if err != nil {
diff --git a/site/layout/page.html b/site/layout/page.html
index 25b733c..62d8b51 100644
--- a/site/layout/page.html
+++ b/site/layout/page.html
@@ -43,4 +43,4 @@ <h1>{{ .Frontmatter.Title }}</h1>
 
 </html>
 
-{{ end}}
+{{ end}}
\ No newline at end of file
diff --git a/site/layout/partials/head.html b/site/layout/partials/head.html
index 99b68b0..7472d5e 100644
--- a/site/layout/partials/head.html
+++ b/site/layout/partials/head.html
@@ -13,12 +13,28 @@
     <link rel="stylesheet" href="https://unpkg.com/highlightjs-copy/dist/highlightjs-copy.min.css" />
     <link rel="alternate" type="application/atom+xml" title="feed" href="/feed.atom" />
 
+
+    {{ if .LiveReload }}
+    <script>
+        const eventSource = new EventSource('http://localhost:8000/events');
+        eventSource.onmessage = function (event) {
+            location.reload()
+        };
+        window.onbeforeunload = function(event)
+        {
+               return confirm("Confirm refresh");
+        };
+    </script>
+    {{ end }}
+
+
     {{range .Frontmatter.JSFiles}}
     <script src="/script/{{.}}" defer></script>
     {{end}} {{range .Layout.SiteScripts}}
     <script src="/script/{{.}}" defer></script>
     {{end}}
 
+
     <meta property="og:type" content="website" />
     <meta property="og:title" content="{{ .Frontmatter.Title }}" />
     <meta property="og:description" content="{{ .Frontmatter.Description }}" />
@@ -33,4 +49,4 @@
 </head>
 {{end}}
 
-</html>
\ No newline at end of file
+</html>