main.go

package main

/*   ____  __
 ___/ / /_/ /____
/ _  / __/ __/ _ \
\_,_/\__/\__/ .__/
           /*/

import (
	"bytes"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"sync"

	_ "embed"

	flag "github.com/spf13/pflag"

	"github.com/zbrd/usage"
)

// Globals
// -------

var (
	empty bytes.Reader
	prog  usage.Program

	version = "v0.0.0-dev"
	data    = []byte("data:\n\n")

	opts = &Options{
		Addr: ":8000",
		Path: ".",
	}

	bus = &Bus{
		su: make(map[int]Sub),
	}

	info = usage.Data{
		"Author":  "Zaim B.",
		"Contact": "zbrd@duck.com",
		"Version": version,
	}
)

//go:embed refresh.js
var refreshJs []byte

//go:embed usage.txt
var usageTpl string

//go:embed version.txt
var versionTpl string

// Types
// -----

// Program options
type Options struct {
	Addr        string
	Path        string
	DoRefresh   bool
	ShowUsage   bool
	ShowVersion bool
}

// Sub is a subscription to a Bus.
type Sub struct {
	Id     int
	Events chan struct{}
}

// Bus is a single-topic event bus.
type Bus struct {
	id int
	su map[int]Sub
	mu sync.RWMutex
}

func (b *Bus) Sub() Sub {
	b.mu.Lock()
	defer b.mu.Unlock()

	sub := Sub{
		Id:     b.id + 1,
		Events: make(chan struct{}),
	}

	b.id = sub.Id
	b.su[b.id] = sub

	slog.Debug("Bus.Sub()", "id", sub.Id)

	return sub
}

func (b *Bus) Unsub(sub Sub) {
	b.mu.Lock()
	defer b.mu.Unlock()

	slog.Debug("Bus.Unsub()", "id", sub.Id)

	close(sub.Events)
	delete(b.su, sub.Id)
	sub.Events = nil
}

func (b *Bus) Pub() {
	b.mu.RLock()
	defer b.mu.RUnlock()

	slog.Debug("Bus.Pub()", "count", len(b.su))

	for _, sub := range b.su {
		select {
		case sub.Events <- struct{}{}:
		}
	}
}

// Main
// ----

func init() {
	prog = usage.Prog(flag.CommandLine)

	flag.Usage = showUsage
	flag.CommandLine.SetOutput(os.Stdout)

	flag.StringVarP(&opts.Addr, "addr", "a", opts.Addr, "")
	flag.StringVarP(&opts.Path, "dir", "d", opts.Path, "")
	flag.BoolVarP(&opts.DoRefresh, "refresh", "R", false, "")
	flag.BoolVarP(&opts.ShowUsage, "help", "h", false, "")
	flag.BoolVarP(&opts.ShowVersion, "version", "v", false, "")

	flag.Lookup("addr").Usage = "Listen at `ADDR`"
	flag.Lookup("dir").Usage = "Serve directory `DIR`"
	flag.Lookup("refresh").Usage = "Send refresh event"
	flag.Lookup("help").Usage = "Show help information"
	flag.Lookup("version").Usage = "Show version information"
}

func main() {
	flag.Parse()

	switch {
	case opts.ShowUsage:
		showUsage()
		return
	case opts.ShowVersion:
		showVersion()
		return
	case opts.DoRefresh:
		doRefresh()
		return
	}

	dir := http.Dir(opts.Path)

	mux := http.NewServeMux()
	mux.Handle("/", http.FileServer(dir))
	mux.HandleFunc("/_/events", handleEvents)
	mux.HandleFunc("/_/refresh.js", handleRefreshJS)
	mux.HandleFunc("POST /_/refresh", handleRefresh)

	h := withLogger(mux)

	slog.Info(
		"Server Listening.",
		"addr", opts.Addr,
		"path", opts.Path,
	)
	if err := http.ListenAndServe(opts.Addr, h); err != nil {
		slog.Error("Failed to listen and serve.", "err", err)
		os.Exit(1)
	}
}

func showUsage() {
	if err := prog.PrintUsage(usageTpl, info); err != nil {
		die("print usage.", err)
	}
}

func showVersion() {
	if err := prog.PrintUsage(versionTpl, info); err != nil {
		die("print version.", err)
	}
}

func doRefresh() {
	var uri string

	if u, err := url.Parse("http://" + opts.Addr); err != nil {
		die("parse address", err)
	} else {
		uri = u.JoinPath("/_/refresh").String()
	}

	if _, err := http.Post(uri, "text/plain", &empty); err != nil {
		die("send refresh event.", err)
	}
}

// Handlers
// --------

func handleEvents(w http.ResponseWriter, r *http.Request) {
	sub := bus.Sub()
	defer bus.Unsub(sub)

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	done := r.Context().Done()
	ctrl := http.NewResponseController(w)

	for {
		select {
		case <-done:
			slog.Debug("handleEvents request done.")
			return
		case <-sub.Events:
			if _, err := w.Write(data); err != nil {
				slog.Error("handleEvents write error.", "err", err)
			}
			if err := ctrl.Flush(); err != nil {
				slog.Error("handleEvents flush error.", "err", err)
			}
		}
	}
}

func handleRefreshJS(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/javascript")
	w.Write(refreshJs)
}

func handleRefresh(w http.ResponseWriter, r *http.Request) {
	bus.Pub()
	w.WriteHeader(204)
}

// Funcs
// -----

func die(op string, err any) {
	fmt.Fprintf(os.Stderr, "Failed to %s: %s", op, err)
	os.Exit(1)
}

func withLogger(h http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		slog.Info(
			"Request",
			"method", r.Method,
			"uri", r.URL.Path,
		)
		h.ServeHTTP(w, r)
	}
}