package healthz

import (
	"context"

	"github.com/cloudwego/hertz/pkg/app"
	"github.com/cloudwego/hertz/pkg/app/server"
)

type Option struct {
	living func() error
	ready  func() error
}
type healthz struct {
	Option
}

func (h *healthz) SetLiving(living func() error) *healthz {
	h.living = living
	return h
}

func (h *healthz) SetReady(ready func() error) *healthz {
	h.ready = ready
	return h
}

func (h *healthz) Living(ctx context.Context, c *app.RequestContext) {
	if h.living == nil {
		c.JSON(200, map[string]string{"ok": "true"})
		return
	}
	if err := h.living(); err != nil {
		c.JSON(500, map[string]string{"ok": "false"})
		return
	}
	c.JSON(200, map[string]string{"ok": "true"})
}
func (h *healthz) Ready(ctx context.Context, c *app.RequestContext) {
	if h.ready == nil {
		c.JSON(200, map[string]string{"ok": "true"})
		return
	}
	if err := h.ready(); err != nil {
		c.JSON(500, map[string]string{"ok": "false"})
		return
	}
	c.JSON(200, map[string]string{"ok": "true"})
}

func (h *healthz) Registry(hertz *server.Hertz) {
	hertz.GET("/living", h.Living)
	hertz.GET("/ready", h.Ready)
}

type Healthz interface {
	SetLiving(living func() error) *healthz
	SetReady(ready func() error) *healthz
	Living(ctx context.Context, c *app.RequestContext)
	Ready(ctx context.Context, c *app.RequestContext)
	Registry(h *server.Hertz)
}

type OptionFunc func(*Option)

func NewHealthz(opts ...OptionFunc) Healthz {
	var opt = Option{}
	for _, o := range opts {
		o(&opt)
	}
	return &healthz{
		Option: opt,
	}
}

func WithLiving(living func() error) OptionFunc {
	return func(o *Option) {
		o.living = living
	}
}

func WithReady(ready func() error) OptionFunc {
	return func(o *Option) {
		o.ready = ready
	}
}