package healthz

import (
	"fmt"
	"net/http"
	"path"
	"sort"
	"strings"

	"github.com/go-logr/logr"
	"k8s.io/apimachinery/pkg/util/sets"
)

// Handler is an http.Handler that aggregates the results of the given
// checkers to the root path, and supports calling individual checkers on
// subpaths of the name of the checker.
//
// Adding checks on the fly is *not* threadsafe -- use a wrapper.
type Handler struct {
	Checks map[string]Checker
	Log    logr.Logger
}

// checkStatus holds the output of a particular check.
type checkStatus struct {
	name     string
	healthy  bool
	excluded bool
}

func (h *Handler) serveAggregated(resp http.ResponseWriter, req *http.Request) {
	failed := false
	excluded := getExcludedChecks(req)

	parts := make([]checkStatus, 0, len(h.Checks))

	// calculate the results...
	for checkName, check := range h.Checks {
		// no-op the check if we've specified we want to exclude the check
		if excluded.Has(checkName) {
			excluded.Delete(checkName)
			parts = append(parts, checkStatus{name: checkName, healthy: true, excluded: true})
			continue
		}
		if err := check(req); err != nil {
			h.Log.Error(err, "healthz check failed", "checker", checkName)
			parts = append(parts, checkStatus{name: checkName, healthy: false})
			failed = true
		} else {
			parts = append(parts, checkStatus{name: checkName, healthy: true})
		}
	}

	// ...default a check if none is present...
	if len(h.Checks) == 0 {
		parts = append(parts, checkStatus{name: "ping", healthy: true})
	}

	for _, c := range excluded.UnsortedList() {
		h.Log.Error(nil, "cannot exclude health check, no matches for it", "checker", c)
	}

	// ...sort to be consistent...
	sort.Slice(parts, func(i, j int) bool { return parts[i].name < parts[j].name })

	// ...and write out the result
	// TODO(help_wanted): this should also accept a request for JSON content (via a accept header)
	_, forceVerbose := req.URL.Query()["verbose"]
	h.writeStatusesAsText(resp, parts, excluded, failed, forceVerbose)
}

// writeStatusAsText writes out the given check statuses in some semi-arbitrary
// bespoke text format that we copied from Kubernetes.  unknownExcludes lists
// any checks that the user requested to have excluded, but weren't actually
// known checks.  writeStatusAsText is always verbose on failure, and can be
// forced to be verbose on success using the given argument.
func (h *Handler) writeStatusesAsText(resp http.ResponseWriter, parts []checkStatus, unknownExcludes sets.Set[string], failed, forceVerbose bool) {
	resp.Header().Set("Content-Type", "text/plain; charset=utf-8")
	resp.Header().Set("X-Content-Type-Options", "nosniff")

	// always write status code first
	if failed {
		resp.WriteHeader(http.StatusInternalServerError)
	} else {
		resp.WriteHeader(http.StatusOK)
	}

	// shortcut for easy non-verbose success
	if !failed && !forceVerbose {
		fmt.Fprint(resp, "ok")
		return
	}

	// we're always verbose on failure, so from this point on we're guaranteed to be verbose

	for _, checkOut := range parts {
		switch {
		case checkOut.excluded:
			fmt.Fprintf(resp, "[+]%s excluded: ok\n", checkOut.name)
		case checkOut.healthy:
			fmt.Fprintf(resp, "[+]%s ok\n", checkOut.name)
		default:
			// don't include the error since this endpoint is public.  If someone wants more detail
			// they should have explicit permission to the detailed checks.
			fmt.Fprintf(resp, "[-]%s failed: reason withheld\n", checkOut.name)
		}
	}

	if unknownExcludes.Len() > 0 {
		fmt.Fprintf(resp, "warn: some health checks cannot be excluded: no matches for %s\n", formatQuoted(unknownExcludes.UnsortedList()...))
	}

	if failed {
		// todo - statuses seems to be broken "statuses":[{},{}],
		h.Log.Error(nil, "healthz check failed", "statuses", parts)
		fmt.Fprintf(resp, "healthz check failed\n")
	} else {
		fmt.Fprint(resp, "healthz check passed\n")
	}
}

func (h *Handler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	// clean up the request (duplicating the internal logic of http.ServeMux a bit)
	// clean up the path a bit
	reqPath := req.URL.Path
	if reqPath == "" || reqPath[0] != '/' {
		reqPath = "/" + reqPath
	}
	// path.Clean removes the trailing slash except for root for us
	// (which is fine, since we're only serving one layer of sub-paths)
	reqPath = path.Clean(reqPath)

	// either serve the root endpoint...
	if reqPath == "/" {
		h.serveAggregated(resp, req)
		return
	}

	// ...the default check (if nothing else is present)...
	if len(h.Checks) == 0 && reqPath[1:] == "ping" {
		CheckHandler{Checker: Ping}.ServeHTTP(resp, req)
		return
	}

	// ...or an individual checker
	checkName := reqPath[1:] // ignore the leading slash
	checker, known := h.Checks[checkName]
	if !known {
		http.NotFoundHandler().ServeHTTP(resp, req)
		return
	}

	CheckHandler{Checker: checker}.ServeHTTP(resp, req)
}

// CheckHandler is an http.Handler that serves a health check endpoint at the root path,
// based on its checker.
type CheckHandler struct {
	Checker
}

func (h CheckHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	if err := h.Checker(req); err != nil {
		http.Error(resp, fmt.Sprintf("internal server error: %v", err), http.StatusInternalServerError)
	} else {
		fmt.Fprint(resp, "ok")
	}
}

// Checker knows how to perform a health check.
type Checker func(req *http.Request) error

// Ping returns true automatically when checked.
var Ping Checker = func(_ *http.Request) error { return nil }

// getExcludedChecks extracts the health check names to be excluded from the query param.
func getExcludedChecks(r *http.Request) sets.Set[string] {
	checks, found := r.URL.Query()["exclude"]
	if found {
		return sets.New[string](checks...)
	}
	return sets.New[string]()
}

// formatQuoted returns a formatted string of the health check names,
// preserving the order passed in.
func formatQuoted(names ...string) string {
	quoted := make([]string, 0, len(names))
	for _, name := range names {
		quoted = append(quoted, fmt.Sprintf("%q", name))
	}
	return strings.Join(quoted, ",")
}