// Package health provides a staleness-aware liveness probe for background-
// loop services (scheduler, execution-worker) that don't otherwise serve HTTP.
//
// The owning loop calls Probe.Bump() on each iteration. A /healthz endpoint
// returns 200 while the last bump is within the staleness window, and 503
// once the loop has gone quiet — so a wedged goroutine (stuck SSH, deadlock)
// surfaces as an unhealthy container instead of a silently-idle one.
package health

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"sync/atomic"
	"time"
)

// Probe tracks the last time the owning loop made progress.
type Probe struct {
	last  atomic.Int64 // unix-nano timestamp of the last Bump
	stale time.Duration
}

// New returns a Probe that considers the owner healthy while Bump has been
// called within stale of the current time.
func New(stale time.Duration) *Probe {
	if stale <= 0 {
		stale = 2 * time.Minute
	}
	p := &Probe{stale: stale}
	p.last.Store(time.Now().UnixNano()) // boot-healthy until first loop stalls
	return p
}

// Bump records that the owning loop completed another iteration.
func (p *Probe) Bump() {
	p.last.Store(time.Now().UnixNano())
}

// Healthy reports whether the last Bump is within the staleness window.
func (p *Probe) Healthy() bool {
	last := time.Unix(0, p.last.Load())
	return time.Since(last) <= p.stale
}

// Handler returns an http.Handler serving GET /healthz. Returns 200 with a
// small JSON body when healthy, 503 (Service Unavailable) when stale.
func (p *Probe) Handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		last := time.Unix(0, p.last.Load())
		if !p.Healthy() {
			w.WriteHeader(http.StatusServiceUnavailable)
		}
		json.NewEncoder(w).Encode(map[string]any{
			"status":         healthStatus(p.Healthy()),
			"last_heartbeat": last.UTC().Format(time.RFC3339),
		})
	})
}

// Serve starts an HTTP server exposing the probe's /healthz on addr until ctx
// is cancelled. A no-op when addr is empty (local/non-docker runs skip it).
// The server is bound to addr (e.g. ":8093"); containers hit it via 127.0.0.1.
func (p *Probe) Serve(ctx context.Context, addr string) {
	if addr == "" {
		return
	}
	mux := http.NewServeMux()
	mux.Handle("/healthz", p.Handler())
	srv := &http.Server{Addr: addr, Handler: mux, ReadHeaderTimeout: 5 * time.Second}

	go func() {
		slog.Info("health server listening", "addr", addr)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			slog.Warn("health server stopped", "addr", addr, "error", err)
		}
	}()
	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		srv.Shutdown(shutdownCtx)
	}()
}

func healthStatus(ok bool) string {
	if ok {
		return "ok"
	}
	return "stale"
}
