using EonaCat.LogStack.Core;
using System;
using System.Threading;

namespace EonaCat.LogStack.Boosters;

/// <summary>
/// Distributed tracing booster that extracts and propagates trace context
/// </summary>
public sealed class DistributedTracingBooster : BoosterBase
{
    private readonly AsyncLocal<string> _correlationId = new AsyncLocal<string>();
    private readonly AsyncLocal<string> _parentSpanId = new AsyncLocal<string>();
    private int _spanIdCounter;

    public DistributedTracingBooster() : base("DistributedTracing") { }

    /// <summary>
    /// Sets the correlation ID for the current async context
    /// </summary>
    public void SetCorrelationId(string correlationId)
    {
        _correlationId.Value = correlationId;
    }

    /// <summary>
    /// Gets the current correlation ID
    /// </summary>
    public string? GetCorrelationId()
    {
        return _correlationId.Value;
    }

    public override bool Boost(ref LogEventBuilder builder)
    {
        var correlationId = _correlationId.Value;
        if (string.IsNullOrEmpty(correlationId))
        {
            correlationId = Guid.NewGuid().ToString("N");
            _correlationId.Value = correlationId;
        }

        builder.WithProperty("correlation_id", correlationId);

        // Generate span ID
        int spanId = Interlocked.Increment(ref _spanIdCounter);
        builder.WithProperty("span_id", spanId.ToString("X8"));

        return true;
    }
}
