// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

namespace System.Diagnostics;

internal delegate ActivitySamplingResult DsesSampleActivityFunc(
    bool hasActivityContext,
    ref ActivityCreationOptions<ActivityContext> options);

internal static class DsesSamplerBuilder
{
    public static DsesSampleActivityFunc CreateParentRatioSampler(double ratio)
    {
        long idUpperBound = ratio <= 0.0
            ? long.MinValue
            : ratio >= 1.0
                ? long.MaxValue
                : (long)(ratio * long.MaxValue);

        return (bool hasActivityContext, ref ActivityCreationOptions<ActivityContext> options) =>
        {
            if (hasActivityContext && options.TraceId != default)
            {
                ActivityContext parentContext = options.Parent;

                ActivitySamplingResult samplingDecision = ParentRatioSampler(idUpperBound, in parentContext, options.TraceId);

                return samplingDecision == ActivitySamplingResult.None
                    && (parentContext == default || parentContext.IsRemote)
                    ? ActivitySamplingResult.PropagationData // If it is the root span or the parent is remote select PropagationData so the trace ID is preserved
                                                             // even if no activity of the trace is recorded
                    : samplingDecision;
            }

            return ActivitySamplingResult.None;
        };
    }

    public static unsafe ActivitySamplingResult ParentRatioSampler(long idUpperBound, in ActivityContext parentContext, ActivityTraceId traceId)
    {
        if (parentContext.TraceId != default)
        {
            return parentContext.TraceFlags.HasFlag(ActivityTraceFlags.Recorded)
                ? ActivitySamplingResult.AllDataAndRecorded
                : ActivitySamplingResult.None;
        }

        Span<byte> traceIdBytes = stackalloc byte[16];
        traceId.CopyTo(traceIdBytes);

        return Math.Abs(GetLowerLong(traceIdBytes)) < idUpperBound
            ? ActivitySamplingResult.AllDataAndRecorded
            : ActivitySamplingResult.None;

        static long GetLowerLong(ReadOnlySpan<byte> bytes)
        {
            long result = 0;
            for (int i = 0; i < 8; i++)
            {
                result <<= 8;
#pragma warning disable CS0675 // Bitwise-or operator used on a sign-extended operand
                result |= bytes[i] & 0xff;
#pragma warning restore CS0675 // Bitwise-or operator used on a sign-extended operand
            }

            return result;
        }
    }

    public static DsesSampleActivityFunc CreateParentRateLimitingSampler(int maximumRatePerSecond)
    {
        Debug.Assert(maximumRatePerSecond > 0, "maximumRatePerSecond must be greater than 0");

        RateLimiter rateLimiter = new RateLimiter(maximumRatePerSecond);

        return (bool hasActivityContext, ref ActivityCreationOptions<ActivityContext> options) =>
        {
            if (hasActivityContext && options.TraceId != default)
            {
                ActivityContext parentContext = options.Parent;

                return parentContext == default || parentContext.IsRemote ? // root or remote activity
                        (rateLimiter.TryAcquire() ? ActivitySamplingResult.AllDataAndRecorded : ActivitySamplingResult.None) :
                        parentContext.TraceFlags.HasFlag(ActivityTraceFlags.Recorded) ? ActivitySamplingResult.AllDataAndRecorded : ActivitySamplingResult.None;
            }

            return ActivitySamplingResult.None;
        };
    }
}