package rpc

import (
	"context"
	"gitee.com/h79/goutils/common/logger"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health/grpc_health_v1"
	"time"
)

type HealthClient struct {
}

func (client *HealthClient) Check(conn *grpc.ClientConn, service string, duration time.Duration) (int, error) {

	c := grpc_health_v1.NewHealthClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*duration)
	defer cancel()

	in := grpc_health_v1.HealthCheckRequest{Service: service}

	rsp, err := c.Check(ctx, &in)
	if err != nil {
		return int(grpc_health_v1.HealthCheckResponse_UNKNOWN), err
	}
	return int(rsp.Status), nil
}

func (client *HealthClient) Watch(conn *grpc.ClientConn, service string, duration time.Duration) (grpc_health_v1.Health_WatchClient, error) {

	c := grpc_health_v1.NewHealthClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*duration)
	defer cancel()

	in := grpc_health_v1.HealthCheckRequest{Service: service}

	rsp, err := c.Watch(ctx, &in)
	if err != nil {
		return nil, err
	}
	return rsp, nil
}

// HealthServer 健康检查实现
type HealthServer struct{}

// Check 实现健康检查接口，这里直接返回健康状态，这里也可以有更复杂的健康检查策略，比如根据服务器负载来返回
func (h *HealthServer) Check(ctx context.Context, req *grpc_health_v1.HealthCheckRequest) (*grpc_health_v1.HealthCheckResponse, error) {
	logger.Debug("RPC:Check, Service: %v", req.Service)
	return &grpc_health_v1.HealthCheckResponse{
		Status: grpc_health_v1.HealthCheckResponse_SERVING,
	}, nil
}

// Watch
// Performs a watch for the serving status of the requested service.
// The server will immediately send back a message indicating the current
// serving status.  It will then subsequently send a new message whenever
// the service's serving status changes.
//
// If the requested service is unknown when the call is received, the
// server will send a message setting the serving status to
// SERVICE_UNKNOWN but will *not* terminate the call.  If at some
// future point, the serving status of the service becomes known, the
// server will send a new message with the service's serving status.
//
// If the call terminates with status UNIMPLEMENTED, then clients
// should assume this method is not supported and should not retry the
// call.  If the call terminates with any other status (including OK),
// clients should retry the call with appropriate exponential backoff.
func (h *HealthServer) Watch(req *grpc_health_v1.HealthCheckRequest, s grpc_health_v1.Health_WatchServer) error {
	logger.Debug("RPC:Watch, Service: %v", req.Service)
	if s != nil {
		rsp := grpc_health_v1.HealthCheckResponse{Status: grpc_health_v1.HealthCheckResponse_SERVING}
		return s.Send(&rsp)
	}
	return nil
}

func HealthCheck(grpcServer *grpc.Server) {
	grpc_health_v1.RegisterHealthServer(grpcServer, &HealthServer{})
}