using System.Text.Json;
using ColaFlow.Modules.Mcp.Application.Services;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Logging;

namespace ColaFlow.Modules.Mcp.Infrastructure.Middleware;

/// <summary>
/// Middleware for authenticating MCP requests using API Keys
/// Only applies to /mcp endpoints
/// </summary>
public class McpApiKeyAuthenticationMiddleware(
    RequestDelegate next,
    ILogger<McpApiKeyAuthenticationMiddleware> logger)
{
    public async Task InvokeAsync(HttpContext context, IMcpApiKeyService apiKeyService)
    {
        // Only apply to /mcp endpoints
        if (!context.Request.Path.StartsWithSegments("/mcp"))
        {
            await next(context);
            return;
        }

        // Extract API Key from Authorization header
        var apiKey = ExtractApiKey(context.Request.Headers);
        if (string.IsNullOrEmpty(apiKey))
        {
            logger.LogWarning("MCP request rejected - Missing API Key");
            await WriteUnauthorizedResponse(context, "Missing API Key. Please provide Authorization: Bearer <api_key> header.");
            return;
        }

        // Get client IP address
        var ipAddress = context.Connection.RemoteIpAddress?.ToString();

        // Validate API Key
        var validationResult = await apiKeyService.ValidateAsync(apiKey, ipAddress, context.RequestAborted);
        if (!validationResult.IsValid)
        {
            logger.LogWarning("MCP request rejected - Invalid API Key: {ErrorMessage}", validationResult.ErrorMessage);
            await WriteUnauthorizedResponse(context, validationResult.ErrorMessage ?? "Invalid API Key");
            return;
        }

        // Set authentication context for downstream handlers
        context.Items["McpAuthType"] = "ApiKey";
        context.Items["McpApiKeyId"] = validationResult.ApiKeyId;
        context.Items["McpTenantId"] = validationResult.TenantId;
        context.Items["McpUserId"] = validationResult.UserId;
        context.Items["McpPermissions"] = validationResult.Permissions;

        logger.LogDebug("MCP request authenticated - ApiKey: {ApiKeyId}, Tenant: {TenantId}, User: {UserId}",
            validationResult.ApiKeyId, validationResult.TenantId, validationResult.UserId);

        await next(context);
    }

    /// <summary>
    /// Extract API Key from Authorization header
    /// Supports: Authorization: Bearer <api_key>
    /// </summary>
    private static string? ExtractApiKey(IHeaderDictionary headers)
    {
        if (!headers.TryGetValue("Authorization", out var authHeader))
        {
            return null;
        }

        var authHeaderValue = authHeader.ToString();
        if (string.IsNullOrWhiteSpace(authHeaderValue))
        {
            return null;
        }

        // Support "Bearer <api_key>" format
        const string bearerPrefix = "Bearer ";
        if (authHeaderValue.StartsWith(bearerPrefix, StringComparison.OrdinalIgnoreCase))
        {
            return authHeaderValue.Substring(bearerPrefix.Length).Trim();
        }

        // Also support direct API key without "Bearer " prefix (for compatibility)
        return authHeaderValue.Trim();
    }

    /// <summary>
    /// Write 401 Unauthorized response in MCP JSON-RPC format
    /// </summary>
    private static async Task WriteUnauthorizedResponse(HttpContext context, string message)
    {
        context.Response.StatusCode = 401;
        context.Response.ContentType = "application/json";

        var errorResponse = new
        {
            jsonrpc = "2.0",
            error = new
            {
                code = -32001, // Custom error code for authentication failure
                message = "Unauthorized",
                data = new { details = message }
            },
            id = (object?)null // We don't have request id yet
        };

        var json = JsonSerializer.Serialize(errorResponse, new JsonSerializerOptions
        {
            PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
            DefaultIgnoreCondition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingNull
        });

        await context.Response.WriteAsync(json);
    }
}
