main.go

package main

import (
	"context"
	"crypto/md5"
	"crypto/rand"
	"embed"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"log/slog"
	"math"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/apex/gateway/v2"
	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-lambda-go/lambda"
)

//go:embed static
var static embed.FS

// NextBus describes when the bus is coming
type NextBus struct {
	EstimatedArrival string `json:"EstimatedArrival"`
	Load             string `json:"Load"`
}

// Service is one bus route calling at a stop
type Service struct {
	ServiceNo string  `json:"ServiceNo"`
	NextBus   NextBus `json:"NextBus"`
	NextBus2  NextBus `json:"NextBus2"`
	NextBus3  NextBus `json:"NextBus3"`
}

// SGBusArrivals describes the response from the datamall API
type SGBusArrivals struct {
	BusStopCode string    `json:"BusStopCode"`
	Services    []Service `json:"Services"`
}

// A fictional bus stop that always has buses coming. It never reaches
// datamall, so the render can be checked - by hand, by the tests, or by the
// deploy - without depending on LTA being up or on an ACCOUNTKEY. It is
// marked noindex and named as a test stop so nobody mistakes it for live data.
const (
	testStopCode = "99999"
	testStopName = "Test Bus Stop"
)

// testArrivals makes up buses relative to now, so the stop always looks live.
// Every offset carries a half minute: on a whole minute the countdown would
// already read one lower by the time the page renders.
func testArrivals(now time.Time) (arrivals SGBusArrivals) {
	sgt := now.In(time.FixedZone("SGT", 8*60*60))
	at := func(d time.Duration) string { return sgt.Add(d).Format(time.RFC3339) }

	arrivals.BusStopCode = testStopCode
	// deliberately not in arrival order, so the sort has something to do
	arrivals.Services = []Service{
		{
			ServiceNo: "666",
			NextBus:   NextBus{EstimatedArrival: at(5*time.Minute + 30*time.Second), Load: "SDA"},
			NextBus2:  NextBus{EstimatedArrival: at(17*time.Minute + 30*time.Second), Load: "LSD"},
			NextBus3:  NextBus{EstimatedArrival: at(32*time.Minute + 30*time.Second), Load: "SEA"},
		},
		{
			ServiceNo: "42",
			NextBus:   NextBus{EstimatedArrival: at(time.Minute + 30*time.Second), Load: "SEA"},
		},
	}
	sortByArrival(arrivals.Services)
	return
}

type Server struct {
	mux      *http.ServeMux
	busStops BusStops
	index    *template.Template
	// now is the clock, swapped out in tests
	now func() time.Time
}

type responseWriter struct {
	http.ResponseWriter
	status      int
	wroteHeader bool
	bytes       int
}

func (rw *responseWriter) WriteHeader(code int) {
	if !rw.wroteHeader {
		rw.status = code
		rw.wroteHeader = true
		rw.ResponseWriter.WriteHeader(code)
	}
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	if !rw.wroteHeader {
		rw.WriteHeader(http.StatusOK)
	}
	n, err := rw.ResponseWriter.Write(b)
	rw.bytes += n
	return n, err
}

func (rw *responseWriter) Status() int {
	return rw.status
}

func (rw *responseWriter) BytesWritten() int {
	return rw.bytes
}

func getLogger(logLevel string) *slog.Logger {
	levelVar := slog.LevelVar{}

	if logLevel != "" {
		if err := levelVar.UnmarshalText([]byte(logLevel)); err != nil {
			panic(fmt.Sprintf("Invalid log level %s: %v", logLevel, err))
		}
	}
	if _, ok := os.LookupEnv("AWS_LAMBDA_FUNCTION_NAME"); ok {
		return slog.New(slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{
			Level: levelVar.Level(),
		}))
	}

	return slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level: levelVar.Level(),
	}))
}

func main() {
	slog.SetDefault(getLogger(os.Getenv("LOGLEVEL")))

	server, err := NewServer("static/all.json")
	if err != nil {
		slog.Error("failed to create server", "error", err)
		os.Exit(1)
	}

	handler := server.middlewareChain(server.mux)

	if _, ok := os.LookupEnv("AWS_LAMBDA_FUNCTION_NAME"); ok {
		slog.Info("starting server", "version", os.Getenv("VERSION"))
		lambda.Start(lambdaHandler(handler))
		return
	} else {
		slog.Info("starting local server", "version", os.Getenv("VERSION"), "port", os.Getenv("PORT"))
		err = http.ListenAndServe(fmt.Sprintf(":%s", os.Getenv("PORT")), handler)
	}

	slog.Error("error listening", "error", err)
}

func NewServer(busStopsPath string) (*Server, error) {
	bs, err := loadBusJSON(busStopsPath)
	if err != nil {
		return nil, fmt.Errorf("load bus stops: %w", err)
	}
	if len(bs) == 0 {
		return nil, fmt.Errorf("bus stop data is empty")
	}

	srv := Server{
		mux:      http.NewServeMux(),
		busStops: bs,
		now:      time.Now,
	}
	funcs := template.FuncMap{
		"version":    func() string { return os.Getenv("VERSION") },
		"totalStops": func() int { return len(srv.busStops) },
		"nameBusStop": func(id string) string {
			if id == testStopCode {
				return testStopName
			}
			return srv.busStops.nameBusStop(id)
		},
		"isTestStop":   func(id string) bool { return id == testStopCode },
		"styleBusStop": styleBusStop,
		"loadClass":    loadClass,
		"countdown":    formatCountdown,
	}
	srv.index, err = template.New("").Funcs(funcs).ParseFS(static, "static/index.html")
	if err != nil {
		return nil, fmt.Errorf("parse index template: %w", err)
	}

	srv.routes()

	return &srv, nil
}

func (s *Server) routes() {
	s.mux.Handle("/{$}", uniqueVisitor(http.HandlerFunc(s.handleIndex)))
	s.mux.HandleFunc("/closest", s.handleClosest)
	s.mux.HandleFunc("/icon", handleIcon)

	fileServer := http.FileServer(http.FS(static))
	s.mux.Handle("/static/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "public, max-age=86400")
		fileServer.ServeHTTP(w, r)
	}))
}

func (s *Server) middlewareChain(handler http.Handler) http.Handler {
	return logRequest(recoverPanic(handler))
}

func recoverPanic(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				slog.Error("recovered from panic", "error", err)
				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			}
		}()
		next.ServeHTTP(w, r)
	})
}

func (s *Server) handleClosest(w http.ResponseWriter, r *http.Request) {
	lat, err := strconv.ParseFloat(r.URL.Query().Get("lat"), 32)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	lng, err := strconv.ParseFloat(r.URL.Query().Get("lng"), 32)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	closest := s.busStops.closest(Point{lat: lat, lng: lng})
	http.Redirect(w, r, fmt.Sprintf("/?id=%s", closest.BusStopCode), http.StatusFound)
}

func (s *Server) handleIndex(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html")

	id := r.URL.Query().Get("id")
	var arriving SGBusArrivals

	if id == testStopCode {
		arriving = testArrivals(s.now())
	} else if id != "" {
		var err error
		arriving, err = busArrivals(id)
		if err != nil {
			slog.Warn("failed to retrieve bus timings", "error", err)
			http.Error(w, fmt.Sprintf("datamall API is returning, %s", err.Error()), http.StatusFailedDependency)
			return
		}
	}

	w.Header().Set("X-Version", os.Getenv("VERSION"))

	data := struct {
		SGBusArrivals
		Now time.Time
	}{arriving, s.now()}
	if err := s.index.ExecuteTemplate(w, "index.html", data); err != nil {
		slog.Error("template failed to render", "error", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// Match app.js: truncate toward zero, using seconds through +/-60 seconds.
func formatCountdown(arrival string, now time.Time) string {
	at, err := time.Parse(time.RFC3339, arrival)
	if err != nil {
		return arrival
	}
	seconds := at.Sub(now).Seconds()
	if math.Abs(seconds) > 60 {
		return fmt.Sprintf("%dm", int(seconds/60))
	}
	return fmt.Sprintf("%ds", int(seconds))
}

func busArrivals(stopID string) (arrivals SGBusArrivals, err error) {
	if stopID == "" {
		return arrivals, fmt.Errorf("invalid stop ID")
	}

	url := fmt.Sprintf("https://datamall2.mytransport.sg/ltaodataservice/v3/BusArrival?BusStopCode=%s", stopID)

	client := &http.Client{
		Timeout: 2 * time.Second,
	}

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return
	}

	req.Header.Add("AccountKey", os.Getenv("ACCOUNTKEY"))

	res, err := client.Do(req)
	if err != nil {
		return
	}

	defer func() {
		if err := res.Body.Close(); err != nil {
			slog.Error("failed to close response body", "error", err)
		}
	}()

	if res.StatusCode != http.StatusOK {
		return arrivals, fmt.Errorf("bad response: %d", res.StatusCode)
	}

	return parseArrivals(res.Body)
}

// parseArrivals decodes a datamall response, soonest bus first
func parseArrivals(r io.Reader) (arrivals SGBusArrivals, err error) {
	err = json.NewDecoder(r).Decode(&arrivals)
	if err != nil {
		slog.Error("failed to decode response", "error", err)
		return
	}

	sortByArrival(arrivals.Services)

	return
}

// sortByArrival puts the buses arriving soonest first
func sortByArrival(services []Service) {
	sort.Slice(services, func(i, j int) bool {
		return services[i].NextBus.EstimatedArrival < services[j].NextBus.EstimatedArrival
	})
}

type Point struct {
	lat float64
	lng float64
}

// BusStop describes a Singaporean (LTA) bus stop
type BusStop struct {
	BusStopCode string  `json:"BusStopCode"`
	RoadName    string  `json:"RoadName"`
	Description string  `json:"Description"`
	Latitude    float64 `json:"Latitude"`
	Longitude   float64 `json:"Longitude"`
}

type BusStops []BusStop

func loadBusJSON(jsonfile string) (bs BusStops, err error) {
	content, err := static.ReadFile(jsonfile)
	if err != nil {
		return
	}
	err = json.Unmarshal(content, &bs)
	if err != nil {
		return
	}

	return
}

func (bs BusStops) closest(location Point) BusStop {
	c := -1
	closestSoFar := math.Inf(1)
	for i := range bs {
		distance := location.distance(Point{bs[i].Latitude, bs[i].Longitude})
		if distance < closestSoFar {
			// Set the return
			c = i
			// Record closest distance
			closestSoFar = distance
		}
	}
	return bs[c]
}

func (bs BusStops) nameBusStop(busStopID string) (description string) {
	for _, p := range bs {
		if busStopID == p.BusStopCode {
			return p.Description
		}
	}
	return ""
}

func styleBusStop(busStopID string) (style template.CSS) {
	data := []byte(busStopID)
	return template.CSS(fmt.Sprintf("background-color: #%.3x; padding: 0.2em", md5.Sum(data)))
}

func loadClass(load string) string {
	switch load {
	case "SEA":
		return "load-seats"
	case "SDA":
		return "load-standing"
	case "LSD":
		return "load-full"
	default:
		return ""
	}
}

func (p Point) distance(p2 Point) float64 {
	latd := p2.lat - p.lat
	lngd := p2.lng - p.lng
	return latd*latd + lngd*lngd
}

func logRequest(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		rw := &responseWriter{
			ResponseWriter: w,
			status:         http.StatusOK,
		}

		defer func() {
			// Log response headers
			for name, values := range rw.Header() {
				for _, value := range values {
					slog.Debug("response header", "name", name, "value", value)
				}
			}

			slog.Info("response",
				"req_method", r.Method,
				"req_ip", r.RemoteAddr,
				"req_path", r.RequestURI,
				"res_status", rw.Status(),
				"res_size", rw.BytesWritten(),
				"res_content_type", rw.Header().Get("Content-Type"),
				"duration", time.Since(start).Milliseconds(),
			)
		}()
		next.ServeHTTP(rw, r)
	})
}

func uniqueVisitor(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cookie, _ := r.Cookie("visitor")
		if cookie == nil || cookie.Value == "" {
			cookie = &http.Cookie{
				Name:    "visitor",
				Value:   "visitor-" + rand.Text(),
				Path:    "/",
				Expires: time.Now().Add(365 * 24 * time.Hour),
			}
			http.SetCookie(w, cookie)
		}
		// Keep the existing log field and message for visitors.sh, including
		// first visits that never load an asset or return to the page.
		slog.Info("return visitor", "unique", cookie.Value)
		next.ServeHTTP(w, r)
	})
}

// lambdaHandler keeps cookies only in the HTTP API v2 cookies field. The
// gateway also copies them into Headers at WriteHeader time, duplicating them.
func lambdaHandler(handler http.Handler) func(context.Context, events.APIGatewayV2HTTPRequest) (events.APIGatewayV2HTTPResponse, error) {
	return func(ctx context.Context, event events.APIGatewayV2HTTPRequest) (events.APIGatewayV2HTTPResponse, error) {
		r, err := gateway.NewRequest(ctx, event)
		if err != nil {
			return events.APIGatewayV2HTTPResponse{}, err
		}
		w := gateway.NewResponse()
		handler.ServeHTTP(w, r)
		response := w.End()
		for key := range response.Headers {
			if strings.EqualFold(key, "Set-Cookie") {
				delete(response.Headers, key)
			}
		}
		for key := range response.MultiValueHeaders {
			if strings.EqualFold(key, "Set-Cookie") {
				delete(response.MultiValueHeaders, key)
			}
		}
		return response, nil
	}
}