package generator

import (
	"bytes"
	"context"
	crand "crypto/rand"
	"encoding/binary"
	"encoding/hex"
	"github.com/go-kratos/kratos/v2/transport"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/trace"
	"io"
	"math/rand"
	"sync"
)

type IDGenerator struct {
	sync.Mutex
	randSource *rand.Rand
}

var _ sdktrace.IDGenerator = &IDGenerator{}

// NewSpanID returns a non-zero span ID from a randomly-chosen sequence.
func (gen *IDGenerator) NewSpanID(ctx context.Context, traceID trace.TraceID) trace.SpanID {
	gen.Lock()
	defer gen.Unlock()
	sid := trace.SpanID{}
	_, _ = gen.randSource.Read(sid[:])
	return sid
}

// NewIDs returns a non-zero trace ID and a non-zero span ID from a
// randomly-chosen sequence.
func (gen *IDGenerator) NewIDs(ctx context.Context) (trace.TraceID, trace.SpanID) {
	gen.Lock()
	defer gen.Unlock()
	var (
		tid       = trace.TraceID{}
		sid       = trace.SpanID{}
		tidReader io.Reader
	)
	tidReader = gen.randSource
	if tr, ok := transport.FromServerContext(ctx); ok {
		traceId := tr.RequestHeader().Get("Trace-Id")
		traceHex, err := hex.DecodeString(traceId)
		if err == nil && traceId != "" {
			tidReader = bytes.NewBuffer(traceHex)
		}
	}

	_, _ = tidReader.Read(tid[:])
	_, _ = gen.randSource.Read(sid[:])

	return tid, sid
}

func DefaultIDGenerator() sdktrace.IDGenerator {
	gen := &IDGenerator{}
	var rngSeed int64
	_ = binary.Read(crand.Reader, binary.LittleEndian, &rngSeed)
	gen.randSource = rand.New(rand.NewSource(rngSeed))
	return gen
}