// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016-present Datadog, Inc.

// Package sampler contains all the logic of the agent-side trace sampling
package sampler

import (
	"math"

	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace"
	"github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace/idx"
	"github.com/DataDog/datadog-agent/pkg/trace/traceutil"
)

const (
	// KeySamplingRateGlobal is a metric key holding the global sampling rate.
	KeySamplingRateGlobal = "_sample_rate"

	// KeySamplingRateClient is a metric key holding the client-set sampling rate for APM events.
	KeySamplingRateClient = "_dd1.sr.rcusr"

	// KeySamplingRatePreSampler is a metric key holding the API rate limiter's rate for APM events.
	KeySamplingRatePreSampler = "_dd1.sr.rapre"

	// KeySamplingRateEventExtraction is the key of the metric storing the event extraction rate on an APM event.
	KeySamplingRateEventExtraction = "_dd1.sr.eausr"

	// KeySamplingRateMaxEPSSampler is the key of the metric storing the max eps sampler rate on an APM event.
	KeySamplingRateMaxEPSSampler = "_dd1.sr.eamax"

	// KeyErrorType is the key of the error type in the meta map
	KeyErrorType = "error.type"

	// KeyAnalyzedSpans is the metric key which specifies if a span is analyzed.
	KeyAnalyzedSpans = "_dd.analyzed"

	// KeyHTTPStatusCode is the key of the http status code in the meta map
	KeyHTTPStatusCode = "http.status_code"

	// KeySpanSamplingMechanism is the metric key holding a span sampling rule that a span was kept on.
	KeySpanSamplingMechanism = "_dd.span_sampling.mechanism"
)

// SamplingPriority is the type encoding a priority sampling decision.
type SamplingPriority int8

const (
	// PriorityNone is the value for SamplingPriority when no priority sampling decision could be found.
	PriorityNone SamplingPriority = math.MinInt8

	// PriorityUserDrop is the value set by a user to explicitly drop a trace.
	PriorityUserDrop SamplingPriority = -1

	// PriorityAutoDrop is the value set by a tracer to suggest dropping a trace.
	PriorityAutoDrop SamplingPriority = 0

	// PriorityAutoKeep is the value set by a tracer to suggest keeping a trace.
	PriorityAutoKeep SamplingPriority = 1

	// PriorityUserKeep is the value set by a user to explicitly keep a trace.
	PriorityUserKeep SamplingPriority = 2

	// 2^64 - 1
	maxTraceID      = ^uint64(0)
	maxTraceIDFloat = float64(maxTraceID)
	// Good number for Knuth hashing (large, prime, fit in int64 for languages without uint64)
	samplerHasher = uint64(1111111111111111111)
)

// IsKeep returns whether the priority is "keep".
func (s SamplingPriority) IsKeep() bool {
	return s == PriorityAutoKeep || s == PriorityUserKeep
}

func (s SamplingPriority) tagValue() string {
	switch s {
	case PriorityUserDrop:
		return "manual_drop"
	case PriorityAutoDrop:
		return "auto_drop"
	case PriorityAutoKeep:
		return "auto_keep"
	case PriorityUserKeep:
		return "manual_keep"
	default:
		return "none"
	}
}

// SampleByRate returns whether to keep a trace, based on its ID and a sampling rate.
// This assumes that trace IDs are nearly uniformly distributed.
func SampleByRate(traceID uint64, rate float64) bool {
	if rate < 1 {
		return traceID*samplerHasher < uint64(rate*maxTraceIDFloat)
	}
	return true
}

// GetSamplingPriority returns the value of the sampling priority metric set on this span and a boolean indicating if
// such a metric was actually found or not.
func GetSamplingPriority(t *pb.TraceChunk) (SamplingPriority, bool) {
	if t.Priority == int32(PriorityNone) {
		return 0, false
	}
	return SamplingPriority(t.Priority), true
}

// GetSamplingPriorityV1 returns the value of the sampling priority metric set on this span and a boolean indicating if
// such a metric was actually found or not.
func GetSamplingPriorityV1(t *idx.InternalTraceChunk) (SamplingPriority, bool) {
	// TODO: would this be PriorityNone?
	if t.Priority == int32(PriorityNone) {
		return 0, false
	}
	return SamplingPriority(t.Priority), true
}

// GetGlobalRate gets the cumulative sample rate of the trace to which this span belongs to.
func GetGlobalRate(s *pb.Span) float64 {
	return getMetricDefault(s, KeySamplingRateGlobal, 1.0)
}

// GetGlobalRateV1 gets the cumulative sample rate of the trace to which this span belongs to.
func GetGlobalRateV1(s *idx.InternalSpan) float64 {
	v, ok := s.GetAttributeAsFloat64(KeySamplingRateGlobal)
	if !ok {
		return 1.0
	}
	return v
}

// GetClientRate gets the rate at which the trace this span belongs to was sampled by the tracer.
// NOTE: This defaults to 1 if no rate is stored.
func GetClientRate(s *pb.Span) float64 {
	return getMetricDefault(s, KeySamplingRateClient, 1.0)
}

// GetClientRateV1 gets the rate at which the trace this span belongs to was sampled by the tracer.
// NOTE: This defaults to 1 if no rate is stored.
func GetClientRateV1(s *idx.InternalSpan) float64 {
	return getMetricDefaultV1(s, KeySamplingRateClient, 1.0)
}

// SetClientRate sets the rate at which the trace this span belongs to was sampled by the tracer.
func SetClientRate(s *pb.Span, rate float64) {
	if rate < 1 {
		setMetric(s, KeySamplingRateClient, rate)
	} else {
		// We assume missing value is 1 to save bandwidth (check getter).
		delete(s.Metrics, KeySamplingRateClient)
	}
}

// SetClientRateV1 sets the rate at which the trace this span belongs to was sampled by the tracer.
func SetClientRateV1(s *idx.InternalSpan, rate float64) {
	if rate < 1 {
		s.SetFloat64Attribute(KeySamplingRateClient, rate)
	} else {
		// We assume missing value is 1 to save bandwidth (check getter).
		s.DeleteAttribute(KeySamplingRateClient)
	}
}

// GetPreSampleRate returns the rate at which the trace this span belongs to was sampled by the agent's presampler.
// NOTE: This defaults to 1 if no rate is stored.
func GetPreSampleRate(s *pb.Span) float64 {
	return getMetricDefault(s, KeySamplingRatePreSampler, 1.0)
}

// GetPreSampleRateV1 returns the rate at which the trace this span belongs to was sampled by the agent's presampler.
// NOTE: This defaults to 1 if no rate is stored.
func GetPreSampleRateV1(s *idx.InternalSpan) float64 {
	return getMetricDefaultV1(s, KeySamplingRatePreSampler, 1.0)
}

// SetPreSampleRate sets the rate at which the trace this span belongs to was sampled by the agent's presampler.
func SetPreSampleRate(s *pb.Span, rate float64) {
	if rate < 1 {
		setMetric(s, KeySamplingRatePreSampler, rate)
	} else {
		// We assume missing value is 1 to save bandwidth (check getter).
		delete(s.Metrics, KeySamplingRatePreSampler)
	}
}

// SetPreSampleRateV1 sets the rate at which the trace this span belongs to was sampled by the agent's presampler.
func SetPreSampleRateV1(s *idx.InternalSpan, rate float64) {
	if rate < 1 {
		s.SetFloat64Attribute(KeySamplingRatePreSampler, rate)
	} else {
		// We assume missing value is 1 to save bandwidth (check getter).
		s.DeleteAttribute(KeySamplingRatePreSampler)
	}
}

// GetEventExtractionRate gets the rate at which the trace from which we extracted this event was sampled at the tracer.
// This defaults to 1 if no rate is stored.
func GetEventExtractionRate(s *pb.Span) float64 {
	return getMetricDefault(s, KeySamplingRateEventExtraction, 1.0)
}

// SetEventExtractionRate sets the rate at which the trace from which we extracted this event was sampled at the tracer.
func SetEventExtractionRate(s *pb.Span, rate float64) {
	if rate < 1 {
		setMetric(s, KeySamplingRateEventExtraction, rate)
	} else {
		// reduce bandwidth, default is assumed 1.0 in backend
		delete(s.Metrics, KeySamplingRateEventExtraction)
	}
}

// SetEventExtractionRateV1 sets the rate at which the trace from which we extracted this event was sampled at the tracer.
func SetEventExtractionRateV1(s *idx.InternalSpan, rate float64) {
	if rate < 1 {
		s.SetFloat64Attribute(KeySamplingRateEventExtraction, rate)
	} else {
		// reduce bandwidth, default is assumed 1.0 in backend
		s.DeleteAttribute(KeySamplingRateEventExtraction)
	}
}

// GetMaxEPSRate gets the rate at which this event was sampled by the max eps event sampler.
func GetMaxEPSRate(s *pb.Span) float64 {
	return getMetricDefault(s, KeySamplingRateMaxEPSSampler, 1.0)
}

// SetMaxEPSRate sets the rate at which this event was sampled by the max eps event sampler.
func SetMaxEPSRate(s *pb.Span, rate float64) {
	if rate < 1 {
		setMetric(s, KeySamplingRateMaxEPSSampler, rate)
	} else {
		// reduce bandwidth, default is assumed 1.0 in backend
		delete(s.Metrics, KeySamplingRateMaxEPSSampler)
	}
}

// SetMaxEPSRateV1 sets the rate at which this event was sampled by the max eps event sampler.
func SetMaxEPSRateV1(s *idx.InternalSpan, rate float64) {
	if rate < 1 {
		s.SetFloat64Attribute(KeySamplingRateMaxEPSSampler, rate)
	} else {
		// reduce bandwidth, default is assumed 1.0 in backend
		s.DeleteAttribute(KeySamplingRateMaxEPSSampler)
	}
}

// SetAnalyzedSpan marks a span analyzed
func SetAnalyzedSpan(s *pb.Span) {
	setMetric(s, KeyAnalyzedSpans, 1)
}

// SetAnalyzedSpanV1 marks a span analyzed
func SetAnalyzedSpanV1(s *idx.InternalSpan) {
	s.SetFloat64Attribute(KeyAnalyzedSpans, 1)
}

// IsAnalyzedSpan checks if a span is analyzed
func IsAnalyzedSpan(s *pb.Span) bool {
	v, _ := getMetric(s, KeyAnalyzedSpans)
	return v == 1
}

// IsAnalyzedSpanV1 checks if a span is analyzed
func IsAnalyzedSpanV1(s *idx.InternalSpan) bool {
	v, _ := s.GetAttributeAsFloat64(KeyAnalyzedSpans)
	return v == 1
}

func weightRoot(s *pb.Span) float32 {
	if s == nil {
		return 1
	}
	clientRate, ok := s.Metrics[KeySamplingRateGlobal]
	if !ok || clientRate <= 0.0 || clientRate > 1.0 {
		clientRate = 1
	}
	preSamplerRate, ok := s.Metrics[KeySamplingRatePreSampler]
	if !ok || preSamplerRate <= 0.0 || preSamplerRate > 1.0 {
		preSamplerRate = 1
	}
	return float32(1.0 / (preSamplerRate * clientRate))
}

func weightRootV1(s *idx.InternalSpan) float32 {
	if s == nil {
		return 1
	}
	clientRate, ok := s.GetAttributeAsFloat64(KeySamplingRateGlobal)
	if !ok || clientRate <= 0.0 || clientRate > 1.0 {
		clientRate = 1
	}
	preSamplerRate, ok := s.GetAttributeAsFloat64(KeySamplingRatePreSampler)
	if !ok || preSamplerRate <= 0.0 || preSamplerRate > 1.0 {
		preSamplerRate = 1
	}
	return float32(1.0 / (preSamplerRate * clientRate))
}

func getMetric(s *pb.Span, k string) (float64, bool) {
	if s.Metrics == nil {
		return 0, false
	}
	val, ok := s.Metrics[k]
	return val, ok
}

// getMetricDefault gets a value in the span Metrics map or default if no value is stored there.
func getMetricDefault(s *pb.Span, k string, def float64) float64 {
	if val, ok := getMetric(s, k); ok {
		return val
	}
	return def
}

// getMetricDefault gets a value in the span Metrics map or default if no value is stored there.
func getMetricDefaultV1(s *idx.InternalSpan, k string, def float64) float64 {
	if val, ok := s.GetAttributeAsFloat64(k); ok {
		return val
	}
	return def
}

// setMetric sets a value in the span Metrics map.
func setMetric(s *pb.Span, key string, val float64) {
	if s.Metrics == nil {
		s.Metrics = make(map[string]float64)
	}
	s.Metrics[key] = val
}

// SingleSpanSampling does single span sampling on the trace, returning true if the trace was modified
func SingleSpanSampling(pt *traceutil.ProcessedTrace) bool {
	ssSpans := getSingleSpanSampledSpans(pt)
	if len(ssSpans) > 0 {
		// Span sampling has kept some spans -> update the chunk
		pt.TraceChunk.Spans = ssSpans
		pt.TraceChunk.Priority = int32(PriorityUserKeep)
		pt.TraceChunk.DroppedTrace = false
		return true
	}
	return false
}

// SingleSpanSamplingV1 does single span sampling on the trace, returning true if the trace was modified
func SingleSpanSamplingV1(pt *traceutil.ProcessedTraceV1) bool {
	ssSpans := getSingleSpanSampledSpansV1(pt)
	if len(ssSpans) > 0 {
		// Span sampling has kept some spans -> update the chunk
		pt.TraceChunk.Spans = ssSpans
		pt.TraceChunk.Priority = int32(PriorityUserKeep)
		pt.TraceChunk.DroppedTrace = false
		return true
	}
	return false
}

// GetSingleSpanSampledSpans searches chunk for spans that have a span sampling tag set and returns them.
func getSingleSpanSampledSpans(pt *traceutil.ProcessedTrace) []*pb.Span {
	var sampledSpans []*pb.Span
	for _, span := range pt.TraceChunk.Spans {
		if _, ok := traceutil.GetMetric(span, KeySpanSamplingMechanism); ok {
			// Keep only those spans that have a span sampling tag.
			sampledSpans = append(sampledSpans, span)
		}
	}
	if sampledSpans == nil {
		// No span sampling tags → no span sampling.
		return nil
	}
	return sampledSpans
}

// GetSingleSpanSampledSpans searches chunk for spans that have a span sampling tag set and returns them.
func getSingleSpanSampledSpansV1(pt *traceutil.ProcessedTraceV1) []*idx.InternalSpan {
	var sampledSpans []*idx.InternalSpan
	for _, span := range pt.TraceChunk.Spans {
		if _, ok := span.GetAttributeAsFloat64(KeySpanSamplingMechanism); ok {
			// Keep only those spans that have a span sampling tag.
			sampledSpans = append(sampledSpans, span)
		}
	}
	if sampledSpans == nil {
		// No span sampling tags → no span sampling.
		return nil
	}
	return sampledSpans
}
