EonaCat.LogStack/EonaCat.LogStack/EonaCatLoggerCore/Tracing/TraceContextManager.cs

using System;
using System.Threading;

namespace EonaCat.LogStack.Tracing
{
    /// <summary>
    /// Manages trace context across async boundaries using AsyncLocal.
    /// Enables automatic trace/span propagation without explicit passing.
    /// </summary>
    public sealed class TraceContextManager
    {
        private static readonly AsyncLocal<TraceContext?> _current = new();
        private static readonly TraceContextManager _instance = new();

        /// <summary>
        /// Gets the singleton instance
        /// </summary>
        public static TraceContextManager Instance => _instance;

        /// <summary>
        /// Gets the current trace context (creates new if not set)
        /// </summary>
        public static TraceContext Current
        {
            get => _current.Value ??= new TraceContext();
            set => _current.Value = value;
        }

        /// <summary>
        /// Gets the current trace context without creating
        /// </summary>
        public static TraceContext? TryGetCurrent() => _current.Value;

        /// <summary>
        /// Sets the trace context
        /// </summary>
        public static void SetCurrent(TraceContext context)
        {
            if (context == null) throw new ArgumentNullException(nameof(context));
            _current.Value = context;
        }

        /// <summary>
        /// Clears the current trace context
        /// </summary>
        public static void Clear()
        {
            _current.Value = null;
        }

        /// <summary>
        /// Creates a new child context and sets it as current
        /// </summary>
        public static IDisposable CreateChildScope()
        {
            var parent = Current;
            var child = parent.CreateChild();
            _current.Value = child;
            return new TraceScope(parent);
        }

        /// <summary>
        /// Executes action with a new child trace context
        /// </summary>
        public static void WithChild(Action action)
        {
            using (CreateChildScope())
            {
                action();
            }
        }

        /// <summary>
        /// Executes function with a new child trace context
        /// </summary>
        public static T WithChild<T>(Func<T> func)
        {
            using (CreateChildScope())
            {
                return func();
            }
        }

        /// <summary>
        /// Gets the current trace ID
        /// </summary>
        public static string GetTraceId() => Current.TraceId;

        /// <summary>
        /// Gets the current span ID
        /// </summary>
        public static string GetSpanId() => Current.SpanId;

        /// <summary>
        /// Gets the current correlation ID
        /// </summary>
        public static string? GetCorrelationId() => Current.CorrelationId;

        /// <summary>
        /// Sets the correlation ID
        /// </summary>
        public static void SetCorrelationId(string correlationId)
        {
            Current.CorrelationId = correlationId;
        }

        /// <summary>
        /// Adds baggage data
        /// </summary>
        public static void AddBaggage(string key, string value)
        {
            Current.Baggage[key] = value;
        }

        /// <summary>
        /// Gets baggage value
        /// </summary>
        public static string? GetBaggage(string key)
        {
            return Current.Baggage.TryGetValue(key, out var value) ? value : null;
        }

        /// <summary>
        /// Removes baggage value
        /// </summary>
        public static bool RemoveBaggage(string key)
        {
            return Current.Baggage.Remove(key);
        }

        /// <summary>
        /// Sets sampling decision
        /// </summary>
        public static void SetSampled(bool sampled)
        {
            Current.TraceFlags = sampled ? (byte)0x01 : (byte)0x00;
        }

        /// <summary>
        /// Gets sampling decision
        /// </summary>
        public static bool IsSampled() => Current.IsSampled;

        /// <summary>
        /// Scope disposed when exiting
        /// </summary>
        private sealed class TraceScope : IDisposable
        {
            private readonly TraceContext _parent;

            public TraceScope(TraceContext parent)
            {
                _parent = parent;
            }

            public void Dispose()
            {
                _current.Value = _parent;
            }
        }
    }
}