using System.Text.Json;
using ColaFlow.Modules.Mcp.Application.Handlers;
using ColaFlow.Modules.Mcp.Contracts.JsonRpc;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Logging;

namespace ColaFlow.Modules.Mcp.Infrastructure.Middleware;

/// <summary>
/// Middleware for handling MCP JSON-RPC 2.0 requests
/// </summary>
public class McpMiddleware
{
    private readonly RequestDelegate _next;
    private readonly ILogger<McpMiddleware> _logger;

    public McpMiddleware(RequestDelegate next, ILogger<McpMiddleware> logger)
    {
        _next = next;
        _logger = logger;
    }

    public async Task InvokeAsync(HttpContext context, IMcpProtocolHandler protocolHandler)
    {
        // Only handle POST requests to /mcp endpoint
        if (context.Request.Method != "POST" || !context.Request.Path.StartsWithSegments("/mcp"))
        {
            await _next(context);
            return;
        }

        _logger.LogDebug("MCP request received from {RemoteIp}", context.Connection.RemoteIpAddress);

        JsonRpcResponse? response = null;
        JsonRpcRequest? request = null;

        try
        {
            // Read request body
            using var reader = new StreamReader(context.Request.Body);
            var requestBody = await reader.ReadToEndAsync();

            _logger.LogTrace("MCP request body: {RequestBody}", requestBody);

            // Parse JSON-RPC request
            try
            {
                request = JsonSerializer.Deserialize<JsonRpcRequest>(requestBody);
                if (request == null)
                {
                    response = JsonRpcResponse.ParseError("Request is null");
                }
            }
            catch (JsonException ex)
            {
                _logger.LogWarning(ex, "Failed to parse JSON-RPC request");
                response = JsonRpcResponse.ParseError(ex.Message);
            }

            // Process request if parsing succeeded
            if (response == null && request != null)
            {
                response = await protocolHandler.HandleRequestAsync(request, context.RequestAborted);
            }

            // Send response (unless it's a notification)
            if (response != null && request?.IsNotification != true)
            {
                context.Response.ContentType = "application/json";
                context.Response.StatusCode = 200;

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

                _logger.LogTrace("MCP response: {ResponseJson}", responseJson);

                await context.Response.WriteAsync(responseJson);
            }
            else if (request?.IsNotification == true)
            {
                // For notifications, return 204 No Content
                context.Response.StatusCode = 204;
                _logger.LogDebug("Notification processed, no response sent");
            }
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Unhandled exception in MCP middleware");

            // Send internal error response (id is null because we don't know the request id)
            response = JsonRpcResponse.InternalError("Unhandled server error", null);
            context.Response.ContentType = "application/json";
            context.Response.StatusCode = 500;

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

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