using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Logging;
using TakeoutSaaS.Shared.Abstractions.Diagnostics;
using TakeoutSaaS.Shared.Abstractions.Ids;

namespace TakeoutSaaS.Shared.Web.Middleware;

/// <summary>
/// 统一 TraceId/CorrelationId，贯穿日志与响应。
/// </summary>
public sealed class CorrelationIdMiddleware(RequestDelegate next, ILogger<CorrelationIdMiddleware> logger, IIdGenerator idGenerator)
{
    private const string TraceHeader = "X-Trace-Id";
    private const string SpanHeader = "X-Span-Id";
    private const string RequestHeader = "X-Request-Id";

    public async Task InvokeAsync(HttpContext context)
    {
        var ownsActivity = Activity.Current is null;
        var activity = Activity.Current ?? new Activity("TakeoutSaaS.Request");

        if (activity.Id is null)
        {
            activity.SetIdFormat(ActivityIdFormat.W3C);
            activity.Start();
        }

        var traceId = activity.TraceId.ToString();
        var spanId = activity.SpanId.ToString();

        if (string.IsNullOrWhiteSpace(traceId))
        {
            traceId = ResolveTraceId(context);
        }

        context.TraceIdentifier = traceId;
        TraceContext.TraceId = traceId;
        TraceContext.SpanId = spanId;

        context.Response.OnStarting(() =>
        {
            context.Response.Headers[TraceHeader] = traceId;
            context.Response.Headers[SpanHeader] = spanId;
            return Task.CompletedTask;
        });

        using (logger.BeginScope(new Dictionary<string, object>
               {
                   ["TraceId"] = traceId,
                   ["SpanId"] = spanId
               }))
        {
            try
            {
                await next(context);
            }
            finally
            {
                TraceContext.Clear();
                if (ownsActivity)
                {
                    activity.Stop();
                }
            }
        }
    }

    private string ResolveTraceId(HttpContext context)
    {
        if (TryGetHeader(context, TraceHeader, out var traceId))
        {
            return traceId;
        }

        if (TryGetHeader(context, RequestHeader, out var requestId))
        {
            return requestId;
        }

        return idGenerator.NextId().ToString();
    }

    private static bool TryGetHeader(HttpContext context, string headerName, out string value)
    {
        if (context.Request.Headers.TryGetValue(headerName, out var values))
        {
            var headerValue = values.ToString();
            if (!string.IsNullOrWhiteSpace(headerValue))
            {
                value = headerValue;
                return true;
            }
        }

        value = string.Empty;
        return false;
    }
}