// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

package exemplar // import "go.opentelemetry.io/otel/sdk/metric/exemplar"

import (
	"context"
	"math"
	"math/rand/v2"
	"sync"
	"sync/atomic"
	"time"

	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/sdk/metric/internal/reservoir"
)

// FixedSizeReservoirProvider returns a provider of [FixedSizeReservoir].
func FixedSizeReservoirProvider(k int) ReservoirProvider {
	return func(attribute.Set) Reservoir {
		return NewFixedSizeReservoir(k)
	}
}

// NewFixedSizeReservoir returns a [FixedSizeReservoir] that samples at most
// k exemplars. If there are k or less measurements made, the Reservoir will
// sample each one. If there are more than k, the Reservoir will then randomly
// sample all additional measurement with a decreasing probability.
func NewFixedSizeReservoir(k int) *FixedSizeReservoir {
	if k < 0 {
		k = 0
	}
	// Use math.MaxInt32 instead of math.MaxUint32 to prevent overflowing int
	// on 32-bit systems.
	if k > math.MaxInt32 {
		k = math.MaxInt32
	}
	return &FixedSizeReservoir{
		storage: newStorage(k),
		// Above we ensure k is positive, and less than MaxInt32.
		nextTracker: newNextTracker(uint32(k)), // nolint: gosec
	}
}

var _ Reservoir = &FixedSizeReservoir{}

// FixedSizeReservoir is a [Reservoir] that samples at most k exemplars. If
// there are k or less measurements made, the Reservoir will sample each one.
// If there are more than k, the Reservoir will then randomly sample all
// additional measurement with a decreasing probability.
type FixedSizeReservoir struct {
	reservoir.ConcurrentSafe
	*storage
	*nextTracker
}

// Offer accepts the parameters associated with a measurement. The
// parameters will be stored as an exemplar if the Reservoir decides to
// sample the measurement.
//
// The passed ctx needs to contain any baggage or span that were active
// when the measurement was made. This information may be used by the
// Reservoir in making a sampling decision.
//
// The time t is the time when the measurement was made. The v and a
// parameters are the value and dropped (filtered) attributes of the
// measurement respectively.
func (r *FixedSizeReservoir) Offer(ctx context.Context, t time.Time, n Value, a []attribute.KeyValue) {
	// The following algorithm is "Algorithm L" from Li, Kim-Hung (4 December
	// 1994). "Reservoir-Sampling Algorithms of Time Complexity
	// O(n(1+log(N/n)))". ACM Transactions on Mathematical Software. 20 (4):
	// 481–493 (https://dl.acm.org/doi/10.1145/198429.198435).
	//
	// A high-level overview of "Algorithm L":
	//   0) Pre-calculate the random count greater than the storage size when
	//      an exemplar will be replaced.
	//   1) Accept all measurements offered until the configured storage size is
	//      reached.
	//   2) Loop:
	//      a) When the pre-calculate count is reached, replace a random
	//         existing exemplar with the offered measurement.
	//      b) Calculate the next random count greater than the existing one
	//         which will replace another exemplars
	//
	// The way a "replacement" count is computed is by looking at `n` number of
	// independent random numbers each corresponding to an offered measurement.
	// Of these numbers the smallest `k` (the same size as the storage
	// capacity) of them are kept as a subset. The maximum value in this
	// subset, called `w` is used to weight another random number generation
	// for the next count that will be considered.
	//
	// By weighting the next count computation like described, it is able to
	// perform a uniformly-weighted sampling algorithm based on the number of
	// samples the reservoir has seen so far. The sampling will "slow down" as
	// more and more samples are offered so as to reduce a bias towards those
	// offered just prior to the end of the collection.
	//
	// This algorithm is preferred because of its balance of simplicity and
	// performance. It will compute three random numbers (the bulk of
	// computation time) for each item that becomes part of the reservoir, but
	// it does not spend any time on items that do not. In particular it has an
	// asymptotic runtime of O(k(1 + log(n/k)) where n is the number of
	// measurements offered and k is the reservoir size.
	//
	// See https://en.wikipedia.org/wiki/Reservoir_sampling for an overview of
	// this and other reservoir sampling algorithms. See
	// https://github.com/MrAlias/reservoir-sampling for a performance
	// comparison of reservoir sampling algorithms.

	count, next := r.incrementCount()
	if count < r.k {
		r.store(ctx, int(count), t, n, a)
	} else if count == next {
		// Overwrite a random existing measurement with the one offered.
		idx := rand.IntN(int(r.k))
		r.store(ctx, idx, t, n, a)
		r.wMu.Lock()
		defer r.wMu.Unlock()
		newCount, newNext := r.loadCountAndNext()
		if newNext < next || newCount < count {
			// This Observe() raced with Collect(), and r.reset() has been
			// called since r.incrementCount(). Skip the call to advance in
			// this case because our exemplar may have been collected in the
			// previous interval.
			return
		}
		r.advance()
	}
}

// Collect returns all the held exemplars.
//
// The Reservoir state is preserved after this call.
func (r *FixedSizeReservoir) Collect(dest *[]Exemplar) {
	r.storage.Collect(dest)
	// Call reset here even though it will reset r.count and restart the random
	// number series. This will persist any old exemplars as long as no new
	// measurements are offered, but it will also prioritize those new
	// measurements that are made over the older collection cycle ones.
	r.reset()
}

func newNextTracker(k uint32) *nextTracker {
	nt := &nextTracker{k: k}
	nt.reset()
	return nt
}

type nextTracker struct {
	// countAndNext holds the current counts in the lower 32 bits and the next
	// value in the upper 32 bits.
	countAndNext atomic.Uint64
	// w is the largest random number in a distribution that is used to compute
	// the next next.
	w float64
	// wMu ensures w is kept consistent with next during advance and reset.
	wMu sync.Mutex
	// k is the number of measurements that can be stored in the reservoir.
	k uint32
}

// reset resets r to the initial state.
func (r *nextTracker) reset() {
	r.wMu.Lock()
	defer r.wMu.Unlock()
	// This resets the number of exemplars known.
	// Random index inserts should only happen after the storage is full.
	r.setCountAndNext(0, r.k)

	// Initial random number in the series used to generate r.next.
	//
	// This is set before r.advance to reset or initialize the random number
	// series. Without doing so it would always be 0 or never restart a new
	// random number series.
	//
	// This maps the uniform random number in (0,1) to a geometric distribution
	// over the same interval. The mean of the distribution is inversely
	// proportional to the storage capacity.
	r.w = math.Exp(math.Log(randomFloat64()) / float64(r.k))

	r.advance()
}

// incrementCount increments the count. It returns the count before the
// increment and the current next value.
func (r *nextTracker) incrementCount() (uint32, uint32) {
	n := r.countAndNext.Add(1)
	// Both count and next are stored in the upper and lower 32 bits, and thus
	// can't overflow.
	return uint32(n&((1<<32)-1) - 1), uint32(n >> 32) // nolint: gosec
}

// incrementNext increments the next value.
func (r *nextTracker) incrementNext(inc uint32) {
	r.countAndNext.Add(uint64(inc) << 32)
}

// setCountAndNext sets the count and next values.
func (r *nextTracker) setCountAndNext(count, next uint32) {
	r.countAndNext.Store(uint64(next)<<32 + uint64(count))
}

func (r *nextTracker) loadCountAndNext() (uint32, uint32) {
	n := r.countAndNext.Load()
	// Both count and next are stored in the upper and lower 32 bits, and thus
	// can't overflow.
	return uint32(n&((1<<32)-1) - 1), uint32(n >> 32) // nolint: gosec
}

// advance updates the count at which the offered measurement will overwrite an
// existing exemplar.
func (r *nextTracker) advance() {
	// Calculate the next value in the random number series.
	//
	// The current value of r.w is based on the max of a distribution of random
	// numbers (i.e. `w = max(u_1,u_2,...,u_k)` for `k` equal to the capacity
	// of the storage and each `u` in the interval (0,w)). To calculate the
	// next r.w we use the fact that when the next exemplar is selected to be
	// included in the storage an existing one will be dropped, and the
	// corresponding random number in the set used to calculate r.w will also
	// be replaced. The replacement random number will also be within (0,w),
	// therefore the next r.w will be based on the same distribution (i.e.
	// `max(u_1,u_2,...,u_k)`). Therefore, we can sample the next r.w by
	// computing the next random number `u` and take r.w as `w * u^(1/k)`.
	r.w *= math.Exp(math.Log(randomFloat64()) / float64(r.k))
	// Use the new random number in the series to calculate the count of the
	// next measurement that will be stored.
	//
	// Given 0 < r.w < 1, each iteration will result in subsequent r.w being
	// smaller. This translates here into the next next being selected against
	// a distribution with a higher mean (i.e. the expected value will increase
	// and replacements become less likely)
	//
	// Important to note, the new r.next will always be at least 1 more than
	// the last r.next.
	r.incrementNext(uint32(math.Log(randomFloat64())/math.Log(1-r.w)) + 1)
}

// randomFloat64 returns, as a float64, a uniform pseudo-random number in the
// open interval (0.0,1.0).
func randomFloat64() float64 {
	// TODO: Use an algorithm that avoids rejection sampling. For example:
	//
	//   const precision = 1 << 53 // 2^53
	//   // Generate an integer in [1, 2^53 - 1]
	//   v := rand.Uint64() % (precision - 1) + 1
	//   return float64(v) / float64(precision)
	f := rand.Float64()
	for f == 0 {
		f = rand.Float64()
	}
	return f
}