using System.Diagnostics;
using Serilog;
using Serilog.Core;
using Serilog.Events;
using Serilog.Sinks.InMemory;
using ZB.MOM.WW.Telemetry;
using ZB.MOM.WW.Telemetry.Serilog;

namespace ZB.MOM.WW.Telemetry.Serilog.Tests;

public sealed class TraceContextEnricherTests
{
    private const string SourceName = "ZB.MOM.WW.Telemetry.Serilog.Tests.TraceContext";

    private static Logger BuildLogger(InMemorySink sink) =>
        new LoggerConfiguration()
            .Enrich.With(new TraceContextEnricher())
            .WriteTo.Sink(sink)
            .CreateLogger();

    private static string? ScalarOrNull(LogEvent logEvent, string propertyName) =>
        logEvent.Properties.TryGetValue(propertyName, out var value) && value is ScalarValue scalar
            ? scalar.Value?.ToString()
            : null;

    [Fact]
    public void Active_activity_stamps_trace_id_and_span_id()
    {
        using var listener = new ActivityListener
        {
            ShouldListenTo = source => source.Name == SourceName,
            Sample = (ref ActivityCreationOptions<ActivityContext> _) =>
                ActivitySamplingResult.AllDataAndRecorded,
        };
        ActivitySource.AddActivityListener(listener);

        using var activitySource = new ActivitySource(SourceName);
        var sink = new InMemorySink();
        using var logger = BuildLogger(sink);

        using var activity = activitySource.StartActivity("unit-test");
        Assert.NotNull(activity);
        Assert.NotNull(Activity.Current);

        // Capture IDs before the log call so assertions are not sensitive to activity
        // lifecycle — Activity.Current may differ after the log call returns.
        var expectedTraceId = activity.TraceId.ToString();
        var expectedSpanId = activity.SpanId.ToString();

        logger.Information("traced");

        var logEvent = Assert.Single(sink.LogEvents);
        Assert.Equal(expectedTraceId, ScalarOrNull(logEvent, "trace_id"));
        Assert.Equal(expectedSpanId, ScalarOrNull(logEvent, "span_id"));
    }

    [Fact]
    public void No_active_activity_omits_trace_id_and_span_id()
    {
        Assert.Null(Activity.Current);

        var sink = new InMemorySink();
        using var logger = BuildLogger(sink);

        logger.Information("untraced");

        var logEvent = Assert.Single(sink.LogEvents);
        Assert.False(logEvent.Properties.ContainsKey("trace_id"));
        Assert.False(logEvent.Properties.ContainsKey("span_id"));
    }
}
