package kk_stage

import (
	"context"
	"gitee.com/cruvie/kk_go_kit/kk_id"
	"github.com/gin-gonic/gin"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
	"go.opentelemetry.io/otel/trace/noop"
	"sync"
)

type Stage struct {
	Ctx         context.Context
	GinCtx      *gin.Context
	Keys        *sync.Map
	Tracer      *Tracer
	TraceId     string
	ServiceName string
	DebugMode   bool
}

func NewStage(ctx context.Context, serviceName string, debugMode bool) (stage *Stage) {
	stage = &Stage{
		Ctx:         ctx,
		Keys:        new(sync.Map),
		Tracer:      NewTracer(serviceName, ctx),
		TraceId:     kk_id.GenerateId().String(),
		ServiceName: serviceName,
		DebugMode:   debugMode,
	}
	return stage
}

// Copy use this to create a new stage in a new goroutine
func (stage *Stage) Copy() *Stage {
	ctx := context.Background()
	newStage := &Stage{
		Ctx:         ctx,
		Keys:        stage.Keys,
		Tracer:      NewTracer(stage.ServiceName, ctx),
		TraceId:     stage.TraceId,
		ServiceName: stage.ServiceName,
		DebugMode:   stage.DebugMode,
	}
	if stage.GinCtx != nil {
		newStage.GinCtx = stage.GinCtx.Copy()
	}
	return newStage
}

func (stage *Stage) SetGinCtx(c *gin.Context) *Stage {
	stage.GinCtx = c
	return stage
}
func (stage *Stage) Set(key any, value any) {
	stage.Keys.Store(key, value)
}

func (stage *Stage) Get(key any) (value any, exists bool) {
	value, ok := stage.Keys.Load(key)
	return value, ok
}
func (stage *Stage) StartTrace(spanName string) (span trace.Span) {
	if !EnableTracer {
		return noop.Span{}
	} else {
		ctx, span := stage.Tracer.Tracer.Start(
			stage.Tracer.Context,
			spanName,
			trace.WithAttributes(attribute.String(TraceIdKey, stage.TraceId)),
		)
		stage.Tracer.Context = ctx
		return span
	}
}