// Package health provides an HTTP server for Kubernetes liveness and readiness probes.
package health

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

// ReadyFunc is called to determine if the service is ready. Return true if ready.
type ReadyFunc func() bool

// Server serves /health and /ready endpoints.
type Server struct {
	port       int
	healthPath string
	readyPath  string
	readyCheck ReadyFunc
	srv        *http.Server
}

// New creates a health server on the given port.
func New(port int, healthPath, readyPath string, readyCheck ReadyFunc) *Server {
	s := &Server{
		port:       port,
		healthPath: healthPath,
		readyPath:  readyPath,
		readyCheck: readyCheck,
	}
	mux := http.NewServeMux()
	mux.HandleFunc(healthPath, s.handleHealth)
	mux.HandleFunc(readyPath, s.handleReady)
	s.srv = &http.Server{
		Addr:              fmt.Sprintf(":%d", port),
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
	}
	return s
}

// Start begins serving in the background. Call Stop to shut down.
func (s *Server) Start() {
	ln, err := net.Listen("tcp", s.srv.Addr)
	if err != nil {
		slog.Error("health server listen failed", "error", err)
		return
	}
	slog.Info("health server started", "port", s.port, "health", s.healthPath, "ready", s.readyPath)
	go func() {
		if err := s.srv.Serve(ln); err != nil && err != http.ErrServerClosed {
			slog.Error("health server error", "error", err)
		}
	}()
}

// Stop gracefully shuts down the server.
func (s *Server) Stop(ctx context.Context) {
	if s.srv != nil {
		_ = s.srv.Shutdown(ctx)
		slog.Info("health server stopped")
	}
}

func (s *Server) handleHealth(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, http.StatusOK, map[string]string{"status": "healthy"})
}

func (s *Server) handleReady(w http.ResponseWriter, _ *http.Request) {
	if s.readyCheck != nil && !s.readyCheck() {
		writeJSON(w, http.StatusServiceUnavailable, map[string]string{"status": "not ready"})
		return
	}
	writeJSON(w, http.StatusOK, map[string]string{"status": "ready"})
}

func writeJSON(w http.ResponseWriter, status int, data any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(data)
}