package kk_grpc

import (
	"context"
	"gitee.com/cruvie/kk_go_kit/kk_stage"
	"google.golang.org/grpc/metadata"
	"log/slog"
	"time"
)

// GetTraceIdFromGrpcCtx
// get traceId from metadata set by
//
//	metadata.AppendToOutgoingContext() or metadata.NewOutgoingContext()
func GetTraceIdFromGrpcCtx(ctx context.Context) (traceId string) {
	md, ok := metadata.FromIncomingContext(ctx)
	if ok {
		strings := md.Get(kk_stage.TraceIdKey)
		if len(strings) > 0 {
			traceId = strings[0]
		}
	} else {
		slog.Error("get TraceId from ctx failed")
	}
	return traceId
}

const TimeoutDur = 10 * time.Second

func NewCallGrpcCtx(stage *kk_stage.Stage) (ctx context.Context, cancelFunc context.CancelFunc) {
	ctx = metadata.AppendToOutgoingContext(stage.Tracer.Context, kk_stage.TraceIdKey, stage.TraceId)
	withTimeoutCtx, cancelFunc := context.WithTimeout(ctx, TimeoutDur)
	return withTimeoutCtx, cancelFunc
}

func NewGrpcStage(tracerCtx context.Context, serviceName string, debugMode bool) (stage *kk_stage.Stage) {
	stage = kk_stage.NewStage(context.Background(), serviceName, debugMode)
	stage.Tracer.Context = tracerCtx
	stage.TraceId = GetTraceIdFromGrpcCtx(tracerCtx)
	return stage
}