using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Logging;
using System.Diagnostics;
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";

    /// <summary>
    /// 管道入口，确保 TraceId/SpanId 贯穿请求。
    /// </summary>
    /// <param name="context">HTTP 上下文。</param>
    public async Task InvokeAsync(HttpContext context)
    {
        // 1. 确保活动存在并启动
        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();
        }

        // 2. 生成/解析 TraceId、SpanId
        var traceId = activity.TraceId.ToString();
        var spanId = activity.SpanId.ToString();

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

        // 3. 写入上下文与响应头
        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;
        });

        // 4. 带 Scope 调用后续中间件
        using (logger.BeginScope(new Dictionary<string, object>
        {
            ["TraceId"] = traceId,
            ["SpanId"] = spanId
        }))
        {
            try
            {
                await next(context);
            }
            finally
            {
                // 5. 清理上下文与活动
                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;
    }
}
