using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Primitives;
using System.Diagnostics;
using System.Threading.Tasks;

namespace StellaOps.Excititor.WebService.Extensions;

internal static class ObservabilityExtensions
{
    private const string TraceHeaderName = "X-Stella-TraceId";
    private const string CorrelationHeaderName = "X-Stella-CorrelationId";
    private const string LegacyCorrelationHeaderName = "X-Correlation-Id";
    private const string CorrelationItemKey = "__stella.correlationId";

    public static IApplicationBuilder UseObservabilityHeaders(this IApplicationBuilder app)
    {
        return app.Use((context, next) =>
        {
            var correlationId = ResolveCorrelationId(context);
            context.Items[CorrelationItemKey] = correlationId;

            context.Response.OnStarting(state =>
            {
                var httpContext = (HttpContext)state;
                ApplyHeaders(httpContext);
                return Task.CompletedTask;
            }, context);

            return next();
        });
    }

    private static void ApplyHeaders(HttpContext context)
    {
        var traceId = Activity.Current?.TraceId.ToString() ?? context.TraceIdentifier;
        if (!string.IsNullOrWhiteSpace(traceId))
        {
            context.Response.Headers[TraceHeaderName] = traceId;
        }

        var correlationId = ResolveCorrelationId(context);
        if (!string.IsNullOrWhiteSpace(correlationId))
        {
            context.Response.Headers[CorrelationHeaderName] = correlationId!;
        }
    }

    private static string ResolveCorrelationId(HttpContext context)
    {
        if (context.Items.TryGetValue(CorrelationItemKey, out var existing) && existing is string cached && !string.IsNullOrWhiteSpace(cached))
        {
            return cached;
        }

        if (TryReadHeader(context.Request.Headers, CorrelationHeaderName, out var headerValue) ||
            TryReadHeader(context.Request.Headers, LegacyCorrelationHeaderName, out headerValue))
        {
            return headerValue!;
        }

        return context.TraceIdentifier;
    }

    private static bool TryReadHeader(IHeaderDictionary headers, string name, out string? value)
    {
        if (headers.TryGetValue(name, out StringValues header) && !StringValues.IsNullOrEmpty(header))
        {
            value = header.ToString();
            return true;
        }

        value = null;
        return false;
    }
}