using System.Text.Json;
using JsonRpcEndpointSample.Rpc;
using JsonRpcEndpointSample.Rpc.Abstractions;
using Microsoft.Extensions.DependencyInjection;

namespace JsonRpcEndpointSample.Mcp;

public sealed class McpDispatcher : IMcpDispatcher
{
    private readonly IServiceProvider _serviceProvider;
    private readonly IEnumerable<RpcToolDescriptor> _toolDescriptors;

    public McpDispatcher(
        IServiceProvider serviceProvider,
        IEnumerable<RpcToolDescriptor> toolDescriptors)
    {
        _serviceProvider = serviceProvider;
        _toolDescriptors = toolDescriptors;
    }

    public async Task<JsonRpcResponse> DispatchAsync(JsonRpcRequest request)
    {
        if (request.JsonRpc != "2.0")
        {
            return JsonRpcResponse.CreateError(request.Id, -32600, "Invalid Request");
        }

        if (string.IsNullOrWhiteSpace(request.Method))
        {
            return JsonRpcResponse.CreateError(request.Id, -32600, "Invalid Request");
        }

        if (request.Method == "initialize")
        {
            return CreateInitializeResponse(request);
        }

        if (request.Method == "tools/list")
        {
            return CreateToolsListResponse(request);
        }

        if (request.Method == "tools/call")
        {
            return await CreateToolsCallResponseAsync(request);
        }

        return JsonRpcResponse.CreateError(request.Id, -32601, "Method not found");
    }

    private static JsonRpcResponse CreateInitializeResponse(JsonRpcRequest request)
    {
        Dictionary<string, object?> toolsCapability = new Dictionary<string, object?>();
        toolsCapability.Add("listChanged", false);

        Dictionary<string, object?> capabilities = new Dictionary<string, object?>();
        capabilities.Add("tools", toolsCapability);

        Dictionary<string, object?> serverInfo = new Dictionary<string, object?>();
        serverInfo.Add("name", "JsonRpcEndpointSample");
        serverInfo.Add("version", "1.0.0");

        Dictionary<string, object?> result = new Dictionary<string, object?>();
        result.Add("protocolVersion", "2025-11-25");
        result.Add("capabilities", capabilities);
        result.Add("serverInfo", serverInfo);

        return JsonRpcResponse.CreateResult(request.Id, result);
    }

    private JsonRpcResponse CreateToolsListResponse(JsonRpcRequest request)
    {
        List<Dictionary<string, object?>> tools = new List<Dictionary<string, object?>>();

        foreach (RpcToolDescriptor descriptor in _toolDescriptors.OrderBy(toolDescriptor => toolDescriptor.Name))
        {
            Dictionary<string, object?> annotations = new Dictionary<string, object?>();
            annotations.Add("readOnlyHint", descriptor.IsReadOnly);

            Dictionary<string, object?> tool = new Dictionary<string, object?>();
            tool.Add("name", descriptor.Name);
            tool.Add("description", descriptor.Description);
            tool.Add("inputSchema", descriptor.InputSchema);
            tool.Add("annotations", annotations);

            tools.Add(tool);
        }

        Dictionary<string, object?> result = new Dictionary<string, object?>();
        result.Add("tools", tools);

        return JsonRpcResponse.CreateResult(request.Id, result);
    }

    private async Task<JsonRpcResponse> CreateToolsCallResponseAsync(JsonRpcRequest request)
    {
        if (request.Params.ValueKind != JsonValueKind.Object)
        {
            return JsonRpcResponse.CreateError(request.Id, -32602, "Invalid params");
        }

        if (!request.Params.TryGetProperty("name", out JsonElement nameElement))
        {
            return JsonRpcResponse.CreateError(request.Id, -32602, "Missing tool name");
        }

        string? toolName = nameElement.GetString();
        if (string.IsNullOrWhiteSpace(toolName))
        {
            return JsonRpcResponse.CreateError(request.Id, -32602, "Missing tool name");
        }

        RpcToolDescriptor? toolDescriptor = _toolDescriptors
            .FirstOrDefault(descriptor => string.Equals(descriptor.Name, toolName, StringComparison.Ordinal));

        if (toolDescriptor == null)
        {
            return JsonRpcResponse.CreateError(request.Id, -32602, "Unknown tool");
        }

        IRpcRequestHandler? handler = _serviceProvider.GetKeyedService<IRpcRequestHandler>(toolName);
        if (handler == null)
        {
            return JsonRpcResponse.CreateError(request.Id, -32602, "Unknown tool");
        }

        JsonElement arguments = default;
        if (request.Params.TryGetProperty("arguments", out JsonElement argumentsElement))
        {
            arguments = argumentsElement;
        }
        else
        {
            using JsonDocument emptyArgumentsDocument = JsonDocument.Parse("{}");
            arguments = emptyArgumentsDocument.RootElement.Clone();
        }

        try
        {
            object? handlerResult = await handler.HandleAsync(arguments);
            return CreateSuccessfulToolCallResponse(request, handlerResult);
        }
        catch (Exception exception)
        {
            return CreateFailedToolCallResponse(request, exception.Message);
        }
    }

    private static JsonRpcResponse CreateSuccessfulToolCallResponse(JsonRpcRequest request, object? handlerResult)
    {
        string text = JsonSerializer.Serialize(handlerResult);

        Dictionary<string, object?> textContent = new Dictionary<string, object?>();
        textContent.Add("type", "text");
        textContent.Add("text", text);

        List<Dictionary<string, object?>> content = new List<Dictionary<string, object?>>();
        content.Add(textContent);

        Dictionary<string, object?> result = new Dictionary<string, object?>();
        result.Add("content", content);
        result.Add("structuredContent", handlerResult);
        result.Add("isError", false);

        return JsonRpcResponse.CreateResult(request.Id, result);
    }

    private static JsonRpcResponse CreateFailedToolCallResponse(JsonRpcRequest request, string message)
    {
        Dictionary<string, object?> textContent = new Dictionary<string, object?>();
        textContent.Add("type", "text");
        textContent.Add("text", message);

        List<Dictionary<string, object?>> content = new List<Dictionary<string, object?>>();
        content.Add(textContent);

        Dictionary<string, object?> result = new Dictionary<string, object?>();
        result.Add("content", content);
        result.Add("isError", true);

        return JsonRpcResponse.CreateResult(request.Id, result);
    }
}
