using Microsoft.AspNetCore.Http;
using Serilog.Context;

namespace ColaFlow.Modules.Mcp.Infrastructure.Middleware;

/// <summary>
/// Middleware that generates or extracts correlation ID for request tracking
/// Adds the correlation ID to:
/// - HttpContext.Items (for access in downstream middleware/handlers)
/// - Response headers (for client-side tracking)
/// - Serilog LogContext (for structured logging)
/// </summary>
public class McpCorrelationIdMiddleware
{
    private readonly RequestDelegate _next;
    private const string CorrelationIdHeaderName = "X-Correlation-Id";

    public McpCorrelationIdMiddleware(RequestDelegate next)
    {
        _next = next;
    }

    public async Task InvokeAsync(HttpContext context)
    {
        // Try to get correlation ID from request header, otherwise generate new one
        var correlationId = GetOrCreateCorrelationId(context);

        // Store in HttpContext.Items for access in downstream middleware/handlers
        context.Items["CorrelationId"] = correlationId;

        // Add to response header so client can track the request
        context.Response.OnStarting(() =>
        {
            if (!context.Response.Headers.ContainsKey(CorrelationIdHeaderName))
            {
                context.Response.Headers.TryAdd(CorrelationIdHeaderName, correlationId);
            }
            return Task.CompletedTask;
        });

        // Add to Serilog LogContext so it appears in all log entries for this request
        using (LogContext.PushProperty("CorrelationId", correlationId))
        {
            await _next(context);
        }
    }

    /// <summary>
    /// Gets correlation ID from request header or generates a new one
    /// </summary>
    private static string GetOrCreateCorrelationId(HttpContext context)
    {
        // Check if client provided a correlation ID
        if (context.Request.Headers.TryGetValue(CorrelationIdHeaderName, out var correlationId)
            && !string.IsNullOrWhiteSpace(correlationId))
        {
            return correlationId.ToString();
        }

        // Generate new correlation ID
        return Guid.NewGuid().ToString("N"); // Use "N" format (no hyphens) for shorter IDs
    }
}
