From 60ed5bda2b3d1c6f2d67287a9d7a434f588cdb14 Mon Sep 17 00:00:00 2001
From: spf13 <steve.francia@gmail.com>
Date: Fri, 16 May 2014 11:48:59 -0400
Subject: [PATCH] Live reloading is working! Still need to incorporate cleanly.

---
 commands/hugo.go                  |   7 ++
 commands/server.go                | 103 +++++++++++++++++++++++++++++-
 docs/layouts/chrome/includes.html |   1 +
 3 files changed, 110 insertions(+), 1 deletion(-)

diff --git a/commands/hugo.go b/commands/hugo.go
index 369403ddb..0a85a300f 100644
--- a/commands/hugo.go
+++ b/commands/hugo.go
@@ -260,6 +260,7 @@ func buildSite(watching ...bool) (err error) {
 	}
 	site.Stats()
 	jww.FEEDBACK.Printf("in %v ms\n", int(1000*time.Since(startTime).Seconds()))
+
 	return nil
 }
 
@@ -321,11 +322,17 @@ func NewWatcher(port int) error {
 				if static_changed {
 					fmt.Print("Static file changed, syncing\n\n")
 					utils.StopOnErr(copyStatic(), fmt.Sprintf("Error copying static files to %s", helpers.AbsPathify(viper.GetString("PublishDir"))))
+
+					// Tell livereload a js file changed to force a hard refresh
+					wsHub.broadcast <- []byte(`{"command":"reload","path":"/x.js","originalPath":"","liveCSS":true}`)
 				}
 
 				if dynamic_changed {
 					fmt.Print("Change detected, rebuilding site\n\n")
 					utils.StopOnErr(buildSite(true))
+
+					// Tell livereload a js file changed to force a hard refresh
+					wsHub.broadcast <- []byte(`{"command":"reload","path":"/x.js","originalPath":"","liveCSS":true}`)
 				}
 			case err := <-watcher.Error:
 				if err != nil {
diff --git a/commands/server.go b/commands/server.go
index 2b2794fe8..4d47bdb9f 100644
--- a/commands/server.go
+++ b/commands/server.go
@@ -14,6 +14,7 @@
 package commands
 
 import (
+	"bytes"
 	"fmt"
 	"net"
 	"net/http"
@@ -21,6 +22,9 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/gorilla/websocket"
+	//"code.google.com/p/go.net/websocket"
+
 	"github.com/spf13/cobra"
 	"github.com/spf13/hugo/helpers"
 	jww "github.com/spf13/jwalterweatherman"
@@ -97,9 +101,106 @@ func serve(port int) {
 
 	fmt.Println("Press ctrl+c to stop")
 
-	err := http.ListenAndServe(":"+strconv.Itoa(port), http.FileServer(http.Dir(helpers.AbsPathify(viper.GetString("PublishDir")))))
+	http.Handle("/", http.FileServer(http.Dir(helpers.AbsPathify(viper.GetString("PublishDir")))))
+	go wsHub.run()
+	http.HandleFunc("/livereload", wsHandler)
+
+	err := http.ListenAndServe(":"+strconv.Itoa(port), nil)
 	if err != nil {
 		jww.ERROR.Printf("Error: %s\n", err.Error())
 		os.Exit(1)
 	}
 }
+
+type hub struct {
+	// Registered connections.
+	connections map[*connection]bool
+
+	// Inbound messages from the connections.
+	broadcast chan []byte
+
+	// Register requests from the connections.
+	register chan *connection
+
+	// Unregister requests from connections.
+	unregister chan *connection
+}
+
+var wsHub = hub{
+	broadcast:   make(chan []byte),
+	register:    make(chan *connection),
+	unregister:  make(chan *connection),
+	connections: make(map[*connection]bool),
+}
+
+func (h *hub) run() {
+	for {
+		select {
+		case c := <-h.register:
+			h.connections[c] = true
+		case c := <-h.unregister:
+			delete(h.connections, c)
+			close(c.send)
+		case m := <-h.broadcast:
+			for c := range h.connections {
+				select {
+				case c.send <- m:
+				default:
+					delete(h.connections, c)
+					close(c.send)
+				}
+			}
+		}
+	}
+}
+
+type connection struct {
+	// The websocket connection.
+	ws *websocket.Conn
+
+	// Buffered channel of outbound messages.
+	send chan []byte
+}
+
+func (c *connection) reader() {
+	for {
+		_, message, err := c.ws.ReadMessage()
+		if err != nil {
+			break
+		}
+		fmt.Println(string(message))
+		switch true {
+		case bytes.Contains(message, []byte(`"command":"hello"`)):
+			wsHub.broadcast <- []byte(`{
+			  "command": "hello",
+			  "protocols": [ "http://livereload.com/protocols/official-7" ],
+			  "serverName": "Hugo"
+			  }`)
+		}
+	}
+	c.ws.Close()
+}
+
+func (c *connection) writer() {
+	for message := range c.send {
+		err := c.ws.WriteMessage(websocket.TextMessage, message)
+		if err != nil {
+			break
+		}
+	}
+	c.ws.Close()
+}
+
+var upgrader = &websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}
+
+func wsHandler(w http.ResponseWriter, r *http.Request) {
+	ws, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		return
+	}
+	c := &connection{send: make(chan []byte, 256), ws: ws}
+	wsHub.register <- c
+	defer func() { wsHub.unregister <- c }()
+	go c.writer()
+	c.reader()
+}
diff --git a/docs/layouts/chrome/includes.html b/docs/layouts/chrome/includes.html
index 51ed6cb97..dbc240d76 100644
--- a/docs/layouts/chrome/includes.html
+++ b/docs/layouts/chrome/includes.html
@@ -3,5 +3,6 @@
     <link href="/static/css/hugofont.css" rel="stylesheet"/>
     <link href='http://fonts.googleapis.com/css?family=Arbutus+Slab&family=Cabin:600&family=Source+Code+Pro' rel='stylesheet' type='text/css'>
     <link rel="stylesheet" href="/static/css/monokai_sublime.css">
+    <script src="/static/js/livereload.js?host=localhost&port=1313"></script>
     <script src="/static/js/highlight.pack.js"></script>
     <script>hljs.initHighlightingOnLoad();</script>