// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Net.WebSockets;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Testing;
using Microsoft.Net.Http.Headers;
using Xunit;

namespace Microsoft.AspNetCore.WebSockets.Test
{
    public class WebSocketCompressionMiddlewareTests : LoggedTest
    {
        [Fact]
        public async Task CompressionNegotiationServerCanChooseSevrverNoContextTakeover()
        {
            await using (var server = KestrelWebSocketHelpers.CreateServer(LoggerFactory, out var port, async context =>
            {
                Assert.True(context.WebSockets.IsWebSocketRequest);
                var webSocket = await context.WebSockets.AcceptWebSocketAsync(new WebSocketAcceptContext()
                {
                    DangerousEnableCompression = true,
                    DisableServerContextTakeover = true
                });
            }))
            {
                using (var client = new HttpClient())
                {
                    var uri = new UriBuilder(new Uri($"ws://127.0.0.1:{port}/"));
                    uri.Scheme = "http";

                    // Craft a valid WebSocket Upgrade request
                    using (var request = new HttpRequestMessage(HttpMethod.Get, uri.ToString()))
                    {
                        SetGenericWebSocketRequest(request);
                        request.Headers.Add(HeaderNames.SecWebSocketExtensions, "permessage-deflate");

                        var response = await client.SendAsync(request);
                        Assert.Equal(HttpStatusCode.SwitchingProtocols, response.StatusCode);
                        Assert.Equal("permessage-deflate; server_no_context_takeover", response.Headers.GetValues(HeaderNames.SecWebSocketExtensions).Aggregate((l, r) => $"{l}; {r}"));
                    }
                }
            }
        }

        [Fact]
        public async Task CompressionNegotiationIgnoredIfNotEnabledOnServer()
        {
            await using (var server = KestrelWebSocketHelpers.CreateServer(LoggerFactory, out var port, async context =>
            {
                Assert.True(context.WebSockets.IsWebSocketRequest);
                var webSocket = await context.WebSockets.AcceptWebSocketAsync();
            }))
            {
                using (var client = new HttpClient())
                {
                    var uri = new UriBuilder(new Uri($"ws://127.0.0.1:{port}/"));
                    uri.Scheme = "http";

                    // Craft a valid WebSocket Upgrade request
                    using (var request = new HttpRequestMessage(HttpMethod.Get, uri.ToString()))
                    {
                        SetGenericWebSocketRequest(request);
                        request.Headers.Add(HeaderNames.SecWebSocketExtensions, "permessage-deflate");

                        var response = await client.SendAsync(request);
                        Assert.Equal(HttpStatusCode.SwitchingProtocols, response.StatusCode);
                        Assert.False(response.Headers.Contains(HeaderNames.SecWebSocketExtensions));
                    }
                }
            }
        }

        [Theory]
        [InlineData("permessage-deflate; server_max_window_bits=14, permessage-deflate; server_max_window_bits=13", "permessage-deflate; server_max_window_bits=13")]
        [InlineData("permessage-deflate; client_max_window_bits=8, permessage-deflate; client_max_window_bits=13", "permessage-deflate; client_max_window_bits=13; server_max_window_bits=13")]
        public async Task CompressionNegotiationCanChooseExtension(string clientHeader, string expectedResponse)
        {
            await using (var server = KestrelWebSocketHelpers.CreateServer(LoggerFactory, out var port, async context =>
            {
                Assert.True(context.WebSockets.IsWebSocketRequest);
                var webSocket = await context.WebSockets.AcceptWebSocketAsync(new WebSocketAcceptContext()
                {
                    DangerousEnableCompression = true,
                    ServerMaxWindowBits = 13
                });
            }))
            {
                using (var client = new HttpClient())
                {
                    var uri = new UriBuilder(new Uri($"ws://127.0.0.1:{port}/"));
                    uri.Scheme = "http";

                    // Craft a valid WebSocket Upgrade request
                    using (var request = new HttpRequestMessage(HttpMethod.Get, uri.ToString()))
                    {
                        SetGenericWebSocketRequest(request);
                        request.Headers.Add(HeaderNames.SecWebSocketExtensions, clientHeader);

                        var response = await client.SendAsync(request);
                        Assert.Equal(HttpStatusCode.SwitchingProtocols, response.StatusCode);
                        Assert.Equal(expectedResponse, response.Headers.GetValues(HeaderNames.SecWebSocketExtensions).Aggregate((l, r) => $"{l}; {r}"));
                    }
                }
            }
        }

        // Smoke test that compression works, we aren't responsible for the specifics of the compression frames
        [Fact]
        public async Task CanSendAndReceiveCompressedData()
        {
            await using (var server = KestrelWebSocketHelpers.CreateServer(LoggerFactory, out var port, async context =>
            {
                Assert.True(context.WebSockets.IsWebSocketRequest);
                using var webSocket = await context.WebSockets.AcceptWebSocketAsync(new WebSocketAcceptContext()
                {
                    DangerousEnableCompression = true,
                    ServerMaxWindowBits = 13
                });

                var serverBuffer = new byte[1024];
                while (true)
                {
                    var result = await webSocket.ReceiveAsync(serverBuffer, CancellationToken.None);
                    if (result.MessageType == WebSocketMessageType.Close)
                    {
                        break;
                    }
                    await webSocket.SendAsync(serverBuffer.AsMemory(0, result.Count), result.MessageType, result.EndOfMessage, default);
                }
                await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, null, default);
            }))
            {
                using (var client = new ClientWebSocket())
                {
                    client.Options.DangerousDeflateOptions = new WebSocketDeflateOptions()
                    {
                        ServerMaxWindowBits = 12,
                        ClientMaxWindowBits = 11,
                    };
                    await client.ConnectAsync(new Uri($"ws://127.0.0.1:{port}/"), CancellationToken.None);
                    var sendCount = 8193;
                    var clientBuf = new byte[sendCount];
                    var receiveBuf = new byte[sendCount];
                    Random.Shared.NextBytes(clientBuf);
                    await client.SendAsync(clientBuf.AsMemory(0, sendCount), WebSocketMessageType.Binary, true, default);
                    var totalRecv = 0;
                    while (totalRecv < sendCount)
                    {
                        var result = await client.ReceiveAsync(receiveBuf.AsMemory(totalRecv), default);
                        totalRecv += result.Count;
                        if (result.EndOfMessage)
                        {
                            Assert.Equal(sendCount, totalRecv);
                            for (var i = 0; i < sendCount; ++i)
                            {
                                Assert.True(clientBuf[i] == receiveBuf[i], $"offset {i} not equal: {clientBuf[i]} == {receiveBuf[i]}");
                            }
                        }
                    }

                    await client.CloseAsync(WebSocketCloseStatus.NormalClosure, null, default);
                }
            }
        }

        private static void SetGenericWebSocketRequest(HttpRequestMessage request)
        {
            request.Headers.Connection.Clear();
            request.Headers.Connection.Add("Upgrade");
            request.Headers.Connection.Add("keep-alive");
            request.Headers.Upgrade.Add(new System.Net.Http.Headers.ProductHeaderValue("websocket"));
            request.Headers.Add(HeaderNames.SecWebSocketVersion, "13");
            // SecWebSocketKey required to be 16 bytes
            request.Headers.Add(HeaderNames.SecWebSocketKey, Convert.ToBase64String(new byte[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 }, Base64FormattingOptions.None));
        }
    }
}
