using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Utils;
using System.ComponentModel;
using System.IO.Pipelines;
using System.Text.Json;
using System.Text.Json.Nodes;

namespace ModelContextProtocol.Tests.Server;

/// <summary>
/// Verifies the server establishes its negotiated protocol version exactly once per stateful session: the
/// initial <see langword="null"/>-to-value transition is allowed and re-sending the same version is an
/// idempotent no-op, but a request that switches to a different (even if otherwise supported) version is
/// rejected with <see cref="McpErrorCode.InvalidRequest"/>. The session is driven over a raw stream
/// transport (the stdio-shaped, stateful path) so the per-request <c>_meta</c> protocol version is fully
/// controlled - the SDK's own client normalizes it on every outgoing request, so only a misbehaving peer
/// can trigger a mid-session change.
/// </summary>
public sealed class NegotiatedProtocolVersionTests : LoggedTest, IAsyncDisposable
{
    private readonly Pipe _clientToServer = new();
    private readonly Pipe _serverToClient = new();
    private readonly CancellationTokenSource _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
    private readonly ServiceProvider _services;
    private readonly McpServer _server;
    private readonly Task _serverTask;
    private readonly StreamWriter _writer;
    private readonly StreamReader _reader;

    public NegotiatedProtocolVersionTests(ITestOutputHelper testOutputHelper)
        : base(testOutputHelper)
    {
        var serviceCollection = new ServiceCollection();
        serviceCollection.AddLogging();
        serviceCollection.AddSingleton<ILoggerProvider>(XunitLoggerProvider);
        serviceCollection.AddSingleton<ILoggerProvider>(MockLoggerProvider);
        serviceCollection
            .AddMcpServer()
            .WithStreamServerTransport(_clientToServer.Reader.AsStream(), _serverToClient.Writer.AsStream())
            .WithTools<EchoTools>();

        _services = serviceCollection.BuildServiceProvider(validateScopes: true);
        _server = _services.GetRequiredService<McpServer>();
        _serverTask = _server.RunAsync(_cts.Token);

        _writer = new StreamWriter(_clientToServer.Writer.AsStream()) { AutoFlush = true };
        _reader = new StreamReader(_serverToClient.Reader.AsStream());
    }

    [Fact]
    public async Task PerRequestProtocolVersion_IsEstablishedOnce_AndRejectsLaterChange()
    {
        var ct = TestContext.Current.CancellationToken;

        // The first request establishes the 2026-07-28 version for the stateful session (null -> 2026-07-28).
        Assert.IsType<JsonRpcResponse>(await RoundTripAsync(id: 1, McpProtocolVersions.July2026ProtocolVersion, ct));

        // Re-sending the same version is an idempotent no-op, not an error.
        Assert.IsType<JsonRpcResponse>(await RoundTripAsync(id: 2, McpProtocolVersions.July2026ProtocolVersion, ct));

        // Switching to a different (still-supported) version mid-session is rejected.
        var error = Assert.IsType<JsonRpcError>(await RoundTripAsync(id: 3, McpProtocolVersions.November2025ProtocolVersion, ct));
        Assert.Equal((int)McpErrorCode.InvalidRequest, error.Error.Code);
        Assert.Contains("protocol version cannot change", error.Error.Message, StringComparison.OrdinalIgnoreCase);

        // The rejected request must not have mutated the negotiated version: the original 2026-07-28 version still works.
        Assert.IsType<JsonRpcResponse>(await RoundTripAsync(id: 4, McpProtocolVersions.July2026ProtocolVersion, ct));
    }

    [Fact]
    public async Task LegacyProtocolVersionMetadata_BeforeInitialize_IsAdvisory()
    {
        var ct = TestContext.Current.CancellationToken;

        Assert.IsType<JsonRpcResponse>(
            await RoundTripAsync(id: 1, McpProtocolVersions.November2025ProtocolVersion, ct));
        Assert.Null(_server.NegotiatedProtocolVersion);
        Assert.DoesNotContain(MockLoggerProvider.LogMessages, m => m.LogLevel >= LogLevel.Warning);

        // The advisory legacy value must not block a subsequent modern request from selecting its era.
        Assert.IsType<JsonRpcResponse>(await RoundTripAsync(id: 2, McpProtocolVersions.July2026ProtocolVersion, ct));
        Assert.Equal(McpProtocolVersions.July2026ProtocolVersion, _server.NegotiatedProtocolVersion);
    }

    [Theory]
    [InlineData("null")]
    [InlineData("{}")]
    [InlineData("[]")]
    [InlineData("42")]
    [InlineData("true")]
    public async Task MalformedProtocolVersionMetadata_BeforeInitialize_IsIgnored(string protocolVersionJson)
    {
        var ct = TestContext.Current.CancellationToken;
        var request = new JsonRpcRequest
        {
            Id = new RequestId(1),
            Method = RequestMethods.ToolsList,
            Params = new JsonObject
            {
                ["_meta"] = new JsonObject
                {
                    [MetaKeys.ProtocolVersion] = JsonNode.Parse(protocolVersionJson),
                },
            },
        };

        // An unreadable value leaves the request on the legacy path without establishing a version.
        Assert.IsType<JsonRpcResponse>(await SendAndReceiveAsync(request, ct));
        Assert.Null(_server.NegotiatedProtocolVersion);

        var warning = Assert.Single(
            MockLoggerProvider.LogMessages,
            m => m.LogLevel == LogLevel.Warning && m.Message.Contains(MetaKeys.ProtocolVersion, StringComparison.Ordinal));
        Assert.Contains(RequestMethods.ToolsList, warning.Message, StringComparison.Ordinal);
    }

    [Fact]
    public async Task MalformedProtocolVersionMetadata_OnServerDiscover_IsRejected()
    {
        var ct = TestContext.Current.CancellationToken;
        var request = new JsonRpcRequest
        {
            Id = new RequestId(1),
            Method = RequestMethods.ServerDiscover,
            Params = new JsonObject
            {
                ["_meta"] = new JsonObject
                {
                    [MetaKeys.ProtocolVersion] = new JsonObject(),
                },
            },
        };

        // server/discover exists only on per-request-metadata revisions, so it cannot fall back to legacy
        // handling. This is what keeps a modern client from silently getting legacy treatment on a transport
        // with no protocol version header.
        var error = Assert.IsType<JsonRpcError>(await SendAndReceiveAsync(request, ct));
        Assert.Equal((int)McpErrorCode.InvalidParams, error.Error.Code);
        Assert.Contains(RequestMethods.ServerDiscover, error.Error.Message, StringComparison.Ordinal);
    }

    [Theory]
    [InlineData(null)]
    [InlineData("null")]
    [InlineData("{}")]
    [InlineData("[]")]
    [InlineData("42")]
    [InlineData("true")]
    public async Task ModernSession_RejectsMissingOrMalformedProtocolVersionMetadata(string? protocolVersionJson)
    {
        var ct = TestContext.Current.CancellationToken;

        Assert.IsType<JsonRpcResponse>(await RoundTripAsync(id: 1, McpProtocolVersions.July2026ProtocolVersion, ct));

        var meta = new JsonObject
        {
            [MetaKeys.ClientCapabilities] = new JsonObject(),
        };
        if (protocolVersionJson is not null)
        {
            meta[MetaKeys.ProtocolVersion] = JsonNode.Parse(protocolVersionJson);
        }

        var request = new JsonRpcRequest
        {
            Id = new RequestId(2),
            Method = RequestMethods.ToolsList,
            Params = new JsonObject
            {
                ["_meta"] = meta,
            },
        };

        // Once the era is established the metadata is required, so an unreadable value is an error rather
        // than a reason to drop back to legacy handling.
        var error = Assert.IsType<JsonRpcError>(await SendAndReceiveAsync(request, ct));
        Assert.Equal((int)McpErrorCode.InvalidParams, error.Error.Code);
        Assert.Contains(MetaKeys.ProtocolVersion, error.Error.Message, StringComparison.Ordinal);
        Assert.Equal(McpProtocolVersions.July2026ProtocolVersion, _server.NegotiatedProtocolVersion);
    }

    [Fact]
    public async Task PerRequestMetadata_ServesRequestMissingClientInfo()
    {
        var ct = TestContext.Current.CancellationToken;

        // clientInfo is optional: a request whose _meta omits it is served.
        var response = await RoundTripAsync(
            id: 1,
            McpProtocolVersions.July2026ProtocolVersion,
            ct,
            includeClientInfo: false);

        Assert.IsType<JsonRpcResponse>(response);
    }

    [Fact]
    public async Task PerRequestMetadata_RejectsRequestMissingClientCapabilities()
    {
        var ct = TestContext.Current.CancellationToken;

        var error = Assert.IsType<JsonRpcError>(
            await RoundTripAsync(
                id: 1,
                McpProtocolVersions.July2026ProtocolVersion,
                ct,
                includeClientCapabilities: false));

        Assert.Equal((int)McpErrorCode.InvalidParams, error.Error.Code);
        Assert.Contains(MetaKeys.ClientCapabilities, error.Error.Message, StringComparison.Ordinal);
        Assert.Null(_server.NegotiatedProtocolVersion);
    }

    [Fact]
    public async Task ServerDiscover_WithoutPerRequestMetadata_IsRejectedBeforeInitialize()
    {
        var ct = TestContext.Current.CancellationToken;

        var request = new JsonRpcRequest
        {
            Id = new RequestId(1),
            Method = RequestMethods.ServerDiscover,
            Params = new JsonObject(),
        };

        var error = Assert.IsType<JsonRpcError>(await SendAndReceiveAsync(request, ct));
        Assert.Equal((int)McpErrorCode.InvalidParams, error.Error.Code);
        Assert.Contains(RequestMethods.ServerDiscover, error.Error.Message, StringComparison.Ordinal);
    }

    [Fact]
    public async Task Initialize_WithPerRequestMetadataProtocolVersion_IsRejected()
    {
        var ct = TestContext.Current.CancellationToken;

        var error = Assert.IsType<JsonRpcError>(
            await RoundTripInitializeAsync(id: 1, McpProtocolVersions.July2026ProtocolVersion, ct));

        Assert.Equal((int)McpErrorCode.UnsupportedProtocolVersion, error.Error.Code);
        Assert.Contains("initialize", error.Error.Message, StringComparison.OrdinalIgnoreCase);
    }

    [Fact]
    public async Task Initialize_IgnoresFutureReservedMetadata()
    {
        var ct = TestContext.Current.CancellationToken;

        var request = new JsonRpcRequest
        {
            Id = new RequestId(1),
            Method = RequestMethods.Initialize,
            Params = JsonSerializer.SerializeToNode(new InitializeRequestParams
            {
                ProtocolVersion = McpProtocolVersions.November2025ProtocolVersion,
                Capabilities = new ClientCapabilities(),
                ClientInfo = new Implementation { Name = "test-client", Version = "1.0.0" },
                Meta = new JsonObject
                {
                    [MetaKeys.ProtocolVersion] = McpProtocolVersions.July2026ProtocolVersion,
                    [MetaKeys.ClientInfo] = "not-an-object",
                    [MetaKeys.ClientCapabilities] = "not-an-object",
                    [MetaKeys.LogLevel] = "not-a-level",
                },
            }, McpJsonUtilities.DefaultOptions),
        };

        Assert.IsType<JsonRpcResponse>(await SendAndReceiveAsync(request, ct));
        Assert.Equal(McpProtocolVersions.November2025ProtocolVersion, _server.NegotiatedProtocolVersion);
        Assert.DoesNotContain(MockLoggerProvider.LogMessages, m => m.LogLevel >= LogLevel.Warning);
    }

    [Fact]
    public async Task LegacySession_IgnoresConflictingKnownLegacyProtocolVersionMetadata()
    {
        var ct = TestContext.Current.CancellationToken;

        Assert.IsType<JsonRpcResponse>(
            await RoundTripInitializeAsync(id: 1, McpProtocolVersions.November2025ProtocolVersion, ct));
        Assert.IsType<JsonRpcResponse>(
            await RoundTripAsync(id: 2, McpProtocolVersions.March2025ProtocolVersion, ct));

        Assert.Equal(McpProtocolVersions.November2025ProtocolVersion, _server.NegotiatedProtocolVersion);
    }

    [Fact]
    public async Task LegacySession_IgnoresModernProtocolVersionMetadata()
    {
        var ct = TestContext.Current.CancellationToken;

        Assert.IsType<JsonRpcResponse>(
            await RoundTripInitializeAsync(id: 1, McpProtocolVersions.November2025ProtocolVersion, ct));

        Assert.IsType<JsonRpcResponse>(
            await RoundTripAsync(id: 2, McpProtocolVersions.July2026ProtocolVersion, ct));
        Assert.Equal(McpProtocolVersions.November2025ProtocolVersion, _server.NegotiatedProtocolVersion);
    }

    [Fact]
    public async Task LegacySession_IgnoresMalformedAuxiliaryMetadata()
    {
        var ct = TestContext.Current.CancellationToken;

        Assert.IsType<JsonRpcResponse>(
            await RoundTripInitializeAsync(id: 1, McpProtocolVersions.November2025ProtocolVersion, ct));

        var request = new JsonRpcRequest
        {
            Id = new RequestId(2),
            Method = RequestMethods.ToolsList,
            Params = new JsonObject
            {
                ["_meta"] = new JsonObject
                {
                    [MetaKeys.ProtocolVersion] = new JsonObject(),
                    [MetaKeys.ClientInfo] = "not-an-object",
                    [MetaKeys.ClientCapabilities] = "not-an-object",
                    [MetaKeys.LogLevel] = "not-a-level",
                },
            },
        };

        Assert.IsType<JsonRpcResponse>(await SendAndReceiveAsync(request, ct));
        Assert.DoesNotContain(MockLoggerProvider.LogMessages, m => m.LogLevel >= LogLevel.Warning);
    }

    [Theory]
    [InlineData(MetaKeys.ClientCapabilities, "\"not-an-object\"")]
    [InlineData(MetaKeys.ClientCapabilities, "null")]
    [InlineData(MetaKeys.ClientCapabilities, "[]")]
    [InlineData(MetaKeys.ClientInfo, "\"not-an-object\"")]
    [InlineData(MetaKeys.LogLevel, "\"not-a-level\"")]
    public async Task ModernRequest_RejectsMalformedMetadata_WithoutEstablishingVersion(string key, string valueJson)
    {
        var ct = TestContext.Current.CancellationToken;
        var meta = PerRequestMetadata();
        meta[key] = JsonNode.Parse(valueJson);
        var request = new JsonRpcRequest
        {
            Id = new RequestId(1),
            Method = RequestMethods.ToolsList,
            Params = new JsonObject
            {
                ["_meta"] = meta,
            },
        };

        var error = Assert.IsType<JsonRpcError>(await SendAndReceiveAsync(request, ct));
        Assert.Equal((int)McpErrorCode.InvalidParams, error.Error.Code);
        Assert.Contains(key, error.Error.Message, StringComparison.Ordinal);
        Assert.Null(_server.NegotiatedProtocolVersion);

        Assert.IsType<JsonRpcResponse>(
            await RoundTripInitializeAsync(id: 2, McpProtocolVersions.November2025ProtocolVersion, ct));
    }

    [Fact]
    public async Task SubscriptionsListen_WithInitializeProtocolVersion_IsRejected()
    {
        var ct = TestContext.Current.CancellationToken;

        Assert.IsType<JsonRpcResponse>(
            await RoundTripInitializeAsync(id: 1, McpProtocolVersions.November2025ProtocolVersion, ct));

        var request = new JsonRpcRequest
        {
            Id = new RequestId(2),
            Method = RequestMethods.SubscriptionsListen,
            Params = new JsonObject(),
        };

        var error = Assert.IsType<JsonRpcError>(await SendAndReceiveAsync(request, ct));
        Assert.Equal((int)McpErrorCode.MethodNotFound, error.Error.Code);
        Assert.Contains(RequestMethods.SubscriptionsListen, error.Error.Message, StringComparison.Ordinal);
    }

    [Theory]
    [InlineData("ping")]
    [InlineData("logging/setLevel")]
    [InlineData("resources/subscribe")]
    [InlineData("resources/unsubscribe")]
    public async Task RemovedMethod_WithPerRequestMetadataProtocolVersion_IsRejected(string method)
    {
        var ct = TestContext.Current.CancellationToken;

        var request = new JsonRpcRequest
        {
            Id = new RequestId(1),
            Method = method,
            Params = new JsonObject
            {
                ["_meta"] = PerRequestMetadata(),
            },
        };

        var error = Assert.IsType<JsonRpcError>(await SendAndReceiveAsync(request, ct));
        Assert.Equal((int)McpErrorCode.MethodNotFound, error.Error.Code);
        Assert.Contains(method, error.Error.Message, StringComparison.Ordinal);
    }

    private async Task<JsonRpcMessage> RoundTripAsync(
        long id,
        string protocolVersion,
        CancellationToken cancellationToken,
        bool includeClientInfo = true,
        bool includeClientCapabilities = true)
    {
        // tools/list is available under both the initialize-handshake and 2026-07-28 revisions (unlike ping/initialize,
        // which the 2026-07-28 protocol removed), so it exercises the version guard rather than the
        // per-method availability gate.
        var meta = new JsonObject
        {
            [MetaKeys.ProtocolVersion] = protocolVersion,
        };

        if (McpProtocolVersions.RequiresPerRequestMetadata(protocolVersion))
        {
            if (includeClientInfo)
            {
                meta[MetaKeys.ClientInfo] = new JsonObject
                {
                    ["name"] = "test-client",
                    ["version"] = "1.0.0",
                };
            }

            if (includeClientCapabilities)
            {
                meta[MetaKeys.ClientCapabilities] = new JsonObject();
            }
        }

        var request = new JsonRpcRequest
        {
            Id = new RequestId(id),
            Method = RequestMethods.ToolsList,
            Params = new JsonObject
            {
                ["_meta"] = meta,
            },
        };

        return await SendAndReceiveAsync(request, cancellationToken);
    }

    private static JsonObject PerRequestMetadata() => new()
    {
        [MetaKeys.ProtocolVersion] = McpProtocolVersions.July2026ProtocolVersion,
        [MetaKeys.ClientInfo] = new JsonObject
        {
            ["name"] = "test-client",
            ["version"] = "1.0.0",
        },
        [MetaKeys.ClientCapabilities] = new JsonObject(),
    };

    private async Task<JsonRpcMessage> RoundTripInitializeAsync(long id, string protocolVersion, CancellationToken cancellationToken)
    {
        var request = new JsonRpcRequest
        {
            Id = new RequestId(id),
            Method = RequestMethods.Initialize,
            Params = JsonSerializer.SerializeToNode(new InitializeRequestParams
            {
                ProtocolVersion = protocolVersion,
                Capabilities = new ClientCapabilities(),
                ClientInfo = new Implementation { Name = "test-client", Version = "1.0.0" },
            }, McpJsonUtilities.DefaultOptions),
        };

        return await SendAndReceiveAsync(request, cancellationToken);
    }

    private async Task<JsonRpcMessage> SendAndReceiveAsync(JsonRpcRequest request, CancellationToken cancellationToken)
    {
        string json = JsonSerializer.Serialize(request, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(JsonRpcMessage)));
#if NET
        await _writer.WriteLineAsync(json.AsMemory(), cancellationToken);
#else
        cancellationToken.ThrowIfCancellationRequested();
        await _writer.WriteLineAsync(json);
#endif

        while (true)
        {
#if NET
            string? line = await _reader.ReadLineAsync(cancellationToken)
                .AsTask()
                .WaitAsync(TestConstants.DefaultTimeout, cancellationToken);
#else
            string? line = await _reader.ReadLineAsync()
                .WaitAsync(TestConstants.DefaultTimeout, cancellationToken);
#endif

            if (line is null)
            {
                throw new InvalidOperationException("Server stream closed before responding.");
            }

            if (string.IsNullOrWhiteSpace(line))
            {
                continue;
            }

            var message = (JsonRpcMessage)JsonSerializer.Deserialize(line, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(JsonRpcMessage)))!;

            // Ignore anything that isn't the response to the request we just sent (e.g. notifications).
            if (message is JsonRpcMessageWithId withId && withId.Id.Equals(request.Id))
            {
                return message;
            }
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _cts.CancelAsync();
        _clientToServer.Writer.Complete();
        _serverToClient.Writer.Complete();

        try
        {
            await _serverTask;
        }
        catch (OperationCanceledException)
        {
        }

        await _services.DisposeAsync();
        _cts.Dispose();
        Dispose();
    }

    [McpServerToolType]
    private sealed class EchoTools
    {
        [McpServerTool, Description("Echoes the input back to the caller.")]
        public static string Echo([Description("The message to echo.")] string message) => message;
    }
}