package grpchealth

import (
	"context"
	"strings"
	"sync/atomic"
	"time"
)

const defaultDependencyBudget = 300 * time.Millisecond

// Dependency 是 readiness 中必须可用的外部依赖。
type Dependency struct {
	Name   string
	Budget time.Duration
	Check  func(ctx context.Context) error
}

// ServingChecker 覆盖无外部硬依赖的占位或 stub gRPC 服务 readiness。
type ServingChecker struct {
	serviceName  string
	dependencies []Dependency
	serving      atomic.Bool
	draining     atomic.Bool
}

// NewServingChecker 创建只依赖 gRPC listener 和 draining 状态的 checker。
func NewServingChecker(serviceName string, dependencies ...Dependency) *ServingChecker {
	return &ServingChecker{serviceName: serviceName, dependencies: dependencies}
}

// ServiceName 返回标准 gRPC health 使用的服务边界名。
func (c *ServingChecker) ServiceName() string {
	return c.serviceName
}

// MarkServing 标记 gRPC listener 已进入 Serve 循环。
func (c *ServingChecker) MarkServing() {
	c.serving.Store(true)
}

// MarkStopped 标记 gRPC listener 已退出 Serve 循环。
func (c *ServingChecker) MarkStopped() {
	c.serving.Store(false)
}

// MarkDraining 让 ready 立即失败，服务下线时不再进入可用池。
func (c *ServingChecker) MarkDraining() {
	c.draining.Store(true)
}

// Ready 返回 listener、draining 和外部硬依赖的聚合 readiness。
func (c *ServingChecker) Ready(ctx context.Context) []Check {
	checks := []Check{
		Bool("grpc_listener", c.serving.Load(), "gRPC listener is not serving"),
		Bool("draining", !c.draining.Load(), "node is draining"),
	}
	for _, dependency := range c.dependencies {
		checks = append(checks, dependency.Ready(ctx))
	}
	return checks
}

// Ready executes a dependency check with a short bounded budget.
func (d Dependency) Ready(ctx context.Context) Check {
	name := strings.TrimSpace(d.Name)
	if name == "" {
		name = "dependency"
	}
	if d.Check == nil {
		return Fail(name, "dependency check is not configured")
	}
	budget := d.Budget
	if budget <= 0 {
		budget = defaultDependencyBudget
	}
	checkCtx, cancel := context.WithTimeout(ctx, budget)
	defer cancel()
	if err := d.Check(checkCtx); err != nil {
		return Fail(name, err.Error())
	}
	return OK(name)
}