package web

import "github.com/gofiber/fiber/v3"

// 注册健康探针
type Probe struct {
	// 存活探针
	livingHook func(ctx fiber.Ctx) error
	// 就绪探针
	readyHook func(ctx fiber.Ctx) error
	// 启动探针
	startHook func(ctx fiber.Ctx) error
}

func NewProbe() *Probe {
	return &Probe{}
}

// 设置存活探针
func (p *Probe) SetLivingProbe(hook func(ctx fiber.Ctx) error) *Probe {
	p.livingHook = hook
	return p
}

// 设置就绪探针
func (p *Probe) SetReadyProbe(hook func(ctx fiber.Ctx) error) *Probe {
	p.readyHook = hook
	return p
}

// 设置启动探针
func (p *Probe) SetStartProbe(hook func(ctx fiber.Ctx) error) *Probe {
	p.startHook = hook
	return p
}

// 默认探针
func defaultProbe(ctx fiber.Ctx) error {
	return ctx.SendStatus(fiber.StatusOK)
}

func (p *Probe) Registry(r fiber.Router) {
	if p.livingHook == nil {
		p.SetLivingProbe(defaultProbe)
	}
	if p.readyHook == nil {
		p.SetReadyProbe(defaultProbe)
	}
	if p.startHook == nil {
		p.SetStartProbe(defaultProbe)
	}

	r.Get("/ready", p.readyHook)
	r.Get("/living", p.livingHook)
	r.Get("/start", p.startHook)
}