workout-suggester/internal/web/server.go

// Package web serves the app: Templ pages, HTMX fragments and an SSE
// endpoint that streams the coach's replies token by token.
package web

import (
	"bytes"
	"context"
	"embed"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/a-h/templ"

	"git.b0b.be/bdeb/workout-suggester/internal/coach"
	"git.b0b.be/bdeb/workout-suggester/internal/config"
	"git.b0b.be/bdeb/workout-suggester/internal/llm"
	"git.b0b.be/bdeb/workout-suggester/internal/routes"
	"git.b0b.be/bdeb/workout-suggester/internal/session"
	"git.b0b.be/bdeb/workout-suggester/internal/weather"
	"git.b0b.be/bdeb/workout-suggester/internal/web/views"
)

//go:embed static
var staticFS embed.FS

const (
	sessionCookie = "ws_sid"
	maxFormLen    = 8 << 10
	routesTimeout = 30 * time.Second
)

// Server is the web app.
type Server struct {
	cfg     config.Config
	llm     *llm.Client
	weather *weather.Client
	routes  *routes.Client
	store   *session.Store
	log     *slog.Logger
}

// New wires a server from config. cacheDir keeps OpenStreetMap lookups
// between runs; empty disables that.
func New(cfg config.Config, cacheDir string, log *slog.Logger) *Server {
	hc := &http.Client{Timeout: 45 * time.Second}
	return &Server{
		cfg: cfg,
		llm: &llm.Client{
			BaseURL: cfg.BaseURL,
			APIKey:  cfg.APIKey,
			Model:   cfg.Model,
			HTTP:    &http.Client{}, // per-request deadlines come from contexts
		},
		weather: &weather.Client{ForecastURL: cfg.ForecastURL, GeocodeURL: cfg.GeocodeURL, HTTP: hc},
		routes:  &routes.Client{URLs: strings.Fields(cfg.OverpassURL), BRouterURL: cfg.BRouterURL, HTTP: hc, CacheDir: cacheDir},
		store:   session.NewStore(),
		log:     log,
	}
}

// Handler returns the app's routes.
func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("GET /static/", http.FileServerFS(staticFS))
	mux.HandleFunc("GET /{$}", s.index)
	mux.HandleFunc("GET /health", s.health)
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { w.Write([]byte("ok")) })
	mux.HandleFunc("POST /plan", s.plan)
	mux.HandleFunc("GET /routes", s.routesCard)
	mux.HandleFunc("GET /gpx/{id}", s.gpx)
	mux.HandleFunc("POST /ask", s.ask)
	mux.HandleFunc("GET /stream/{id}", s.stream)
	return s.withSession(mux)
}

// Run serves until ctx is cancelled, then shuts down gracefully.
func (s *Server) Run(ctx context.Context) error {
	srv := &http.Server{
		Addr:              s.cfg.Addr,
		Handler:           s.Handler(),
		ReadHeaderTimeout: 10 * time.Second,
		// No WriteTimeout: replies stream for as long as the model talks.
		// LLM_TIMEOUT bounds each reply instead.
	}

	go func() {
		t := time.NewTicker(10 * time.Minute)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				s.store.Prune(12 * time.Hour)
			}
		}
	}()

	errc := make(chan error, 1)
	go func() { errc <- srv.ListenAndServe() }()
	s.log.Info("workout-suggester listening", "url", "http://"+s.cfg.Addr, "llm", s.cfg.BaseURL, "model", s.cfg.Model)

	select {
	case err := <-errc:
		return err
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		return srv.Shutdown(shutdownCtx)
	}
}

type ctxKey struct{}

// withSession gives every browser a random session cookie.
func (s *Server) withSession(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sid := ""
		if c, err := r.Cookie(sessionCookie); err == nil && len(c.Value) == 32 {
			sid = c.Value
		} else {
			sid = session.NewID()
			http.SetCookie(w, &http.Cookie{
				Name:     sessionCookie,
				Value:    sid,
				Path:     "/",
				HttpOnly: true,
				SameSite: http.SameSiteLaxMode,
			})
		}
		next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), ctxKey{}, sid)))
	})
}

func sessionID(r *http.Request) string {
	sid, _ := r.Context().Value(ctxKey{}).(string)
	return sid
}

func (s *Server) render(w http.ResponseWriter, r *http.Request, c templ.Component) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	if err := c.Render(r.Context(), w); err != nil {
		s.log.Error("render", "err", err)
	}
}

func (s *Server) index(w http.ResponseWriter, r *http.Request) {
	d := views.PageData{DefaultPlace: s.cfg.DefaultPlace}
	if conv := s.store.Get(sessionID(r)); conv != nil {
		p := conv.Plan()
		d.Plan, d.Messages = &p, conv.Messages()
	}
	s.render(w, r, views.Page(d))
}

func (s *Server) health(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
	defer cancel()
	model, err := s.llm.ResolveModel(ctx)
	if err == nil && s.cfg.Model != "" {
		_, err = s.llm.Models(ctx) // configured model: still check the server is up
	}
	h := views.HealthInfo{OK: err == nil, Model: model}
	if err != nil {
		h.Err = err.Error()
	}
	s.render(w, r, views.Health(h))
}

// parseRequest reads the form into a coach request, with safe defaults.
func parseRequest(r *http.Request) coach.Request {
	req := coach.Request{
		Activity:  coach.Activity(r.FormValue("activity")),
		Intensity: coach.Intensity(r.FormValue("intensity")),
		Note:      strings.TrimSpace(r.FormValue("note")),
	}
	switch req.Activity {
	case coach.Run, coach.Ride, coach.Walk, coach.Indoor:
	default:
		req.Activity = coach.Any
	}
	switch req.Intensity {
	case coach.Easy, coach.Moderate, coach.Hard:
	default:
		req.Intensity = coach.Moderate
	}
	req.Minutes, _ = strconv.Atoi(r.FormValue("minutes"))
	req.Minutes = min(240, max(15, req.Minutes))
	if len(req.Note) > 200 {
		req.Note = req.Note[:200]
	}
	return req
}

// plan checks the weather, scores the options and starts the coach's reply.
func (s *Server) plan(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, maxFormLen)
	if err := r.ParseForm(); err != nil {
		s.render(w, r, views.Error("That form was too big."))
		return
	}
	req := parseRequest(r)
	ctx, cancel := context.WithTimeout(r.Context(), 15*time.Second)
	defer cancel()

	var p coach.Plan
	lat, errLat := strconv.ParseFloat(r.FormValue("lat"), 64)
	lon, errLon := strconv.ParseFloat(r.FormValue("lon"), 64)
	switch place := strings.TrimSpace(r.FormValue("place")); {
	case errLat == nil && errLon == nil && lat >= -90 && lat <= 90 && lon >= -180 && lon <= 180:
		p.At = routes.Point{Lat: lat, Lon: lon}
	case place != "":
		pl, err := s.weather.Geocode(ctx, place)
		if err != nil {
			s.log.Warn("geocode", "err", err)
			s.render(w, r, views.Error("Couldn't find that place. Check the spelling, or tap 📍 Here."))
			return
		}
		p.Place, p.At = pl.Name, routes.Point{Lat: pl.Lat, Lon: pl.Lon}
	default:
		s.render(w, r, views.Error("Where are you? Type a town or tap 📍 Here."))
		return
	}

	f, err := s.weather.Forecast(ctx, p.At.Lat, p.At.Lon)
	if err != nil {
		s.log.Warn("forecast", "err", err)
		s.render(w, r, views.Error("Couldn't get the weather right now. Are you online?"))
		return
	}
	p.Forecast, p.Request, p.Assessment = f, req, coach.Assess(f, req)

	conv, _ := s.store.Start(sessionID(r), p)
	s.render(w, r, views.PlanView(conv.Plan(), conv.Messages()))
}

// loadRoutes looks up routes for the plan's pick once per plan. Concurrent
// callers (the routes card and the first reply) share one Overpass request.
func (s *Server) loadRoutes(ctx context.Context, conv *session.Conversation) coach.Plan {
	p := conv.Plan()
	if p.RoutesLoaded || p.Assessment.Pick == coach.Indoor {
		return p
	}
	ctx, cancel := context.WithTimeout(ctx, routesTimeout)
	defer cancel()

	// Loops from the door (BRouter) and signposted routes (Overpass) at once.
	kind := p.Assessment.Pick.RouteKind()
	var (
		loops    []routes.Suggestion
		loopsErr error
		wg       sync.WaitGroup
	)
	wg.Go(func() {
		loops, loopsErr = s.routes.Loops(ctx, kind, p.At, p.Assessment.TargetKm, p.Forecast.Now.WindDirection)
	})
	area, err := s.routes.Near(ctx, p.At.Lat, p.At.Lon)
	wg.Wait()
	if loopsErr != nil {
		s.log.Warn("loops", "err", loopsErr)
	}
	if err != nil {
		s.log.Warn("routes", "err", err)
	}

	conv.UpdatePlan(func(p *coach.Plan) {
		if p.RoutesLoaded {
			return
		}
		p.RoutesLoaded = true
		p.Routes = loops[:min(2, len(loops))]
		if err == nil {
			p.Routes = append(p.Routes, routes.Suggest(area, kind, p.Assessment.TargetKm, p.At, 4-len(p.Routes))...)
			if j, km, ok := area.NearestJunction(kind, p.At); ok && km < 2 {
				p.Junction = &coach.JunctionInfo{Ref: j.Ref, Kind: kind, Km: km}
			}
		}
		if len(p.Routes) == 0 && err != nil {
			p.RoutesErr = err.Error()
		}
	})
	return conv.Plan()
}

func (s *Server) routesCard(w http.ResponseWriter, r *http.Request) {
	conv := s.store.Get(sessionID(r))
	if conv == nil {
		s.render(w, r, views.Error("This plan expired. Ask again."))
		return
	}
	s.render(w, r, views.RoutesCard(s.loadRoutes(r.Context(), conv)))
}

func (s *Server) gpx(w http.ResponseWriter, r *http.Request) {
	id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
	route, ok := s.routes.Route(id)
	if err != nil || !ok {
		http.NotFound(w, r)
		return
	}
	w.Header().Set("Content-Type", "application/gpx+xml")
	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="osm-%d.gpx"`, id))
	routes.WriteGPX(w, route)
}

func (s *Server) ask(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, maxFormLen)
	text := strings.TrimSpace(r.FormValue("message"))
	conv := s.store.Get(sessionID(r))
	switch {
	case conv == nil:
		s.render(w, r, views.Error("This plan expired. Ask again from the top."))
		return
	case text == "" || len(text) > 1000:
		http.Error(w, "message is empty or too long", http.StatusBadRequest)
		return
	}
	user, reply := conv.Ask(text)
	s.render(w, r, views.Exchange(user, reply))
}

// stream generates a reply and pushes it to the browser as server-sent
// events. EventSource reconnects automatically, so a reply that is already
// claimed or finished is answered with its current state instead of being
// generated twice.
func (s *Server) stream(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	conv := s.store.Get(sessionID(r))

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("X-Accel-Buffering", "no") // don't let reverse proxies buffer
	sse := &sseWriter{w: w, rc: http.NewResponseController(w)}

	lost := session.Message{ID: id, Role: "assistant", State: session.Done,
		Err: "This reply was lost, probably because the server restarted. Please ask again."}
	if conv == nil {
		sse.component(r.Context(), "done", views.MessageView(lost))
		return
	}
	// The first reply should know the routes, so wait for them (cached
	// after the first lookup, and shared with the routes card).
	if conv.IsFirstReply(id) {
		s.loadRoutes(r.Context(), conv)
	}
	prompt, msg, ok := conv.Claim(id)
	if !ok {
		if msg.ID == "" {
			msg = lost
		}
		sse.component(r.Context(), "done", views.MessageView(msg))
		return
	}

	ctx, cancel := context.WithTimeout(r.Context(), s.cfg.Timeout)
	defer cancel()
	start := time.Now()
	err := s.llm.Stream(ctx, prompt, func(delta string) error {
		conv.Append(id, delta)
		return sse.component(ctx, "token", views.Token(delta))
	})

	if r.Context().Err() != nil { // browser went away; keep what we have
		conv.Finish(id, errors.New("interrupted"))
		return
	}
	switch {
	case errors.Is(err, context.DeadlineExceeded):
		err = fmt.Errorf("reply took longer than %s and was cut off", s.cfg.Timeout)
	case err != nil:
		s.log.Warn("generation failed", "err", err)
		err = errors.New("the local model didn't answer. Is your model server running?")
	default:
		s.log.Debug("reply done", "id", id, "took", time.Since(start))
	}
	sse.component(r.Context(), "done", views.MessageView(conv.Finish(id, err)))
}

type sseWriter struct {
	w  http.ResponseWriter
	rc *http.ResponseController
}

// component sends one rendered Templ component as a named SSE event.
func (s *sseWriter) component(ctx context.Context, event string, c templ.Component) error {
	var buf bytes.Buffer
	if err := c.Render(ctx, &buf); err != nil {
		return err
	}
	var b strings.Builder
	b.WriteString("event: " + event + "\n")
	for _, line := range strings.Split(buf.String(), "\n") {
		b.WriteString("data: " + line + "\n")
	}
	b.WriteString("\n")
	if _, err := s.w.Write([]byte(b.String())); err != nil {
		return err
	}
	return s.rc.Flush()
}