// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using Azure.Iot.Operations.Protocol.Events;
using Azure.Iot.Operations.Protocol.Models;
using Azure.Iot.Operations.Mqtt.Session;
using Azure.Iot.Operations.Protocol.RPC;
using TestEnvoys.Counter;
using System.Text.Json;

namespace Azure.Iot.Operations.Protocol.IntegrationTests;

public class CounterEnvoyTests
{
    [Fact]
    public async Task IncrementTest()
    {
        ApplicationContext applicationContext = new ApplicationContext();
        string executorId = "counter-server-" + Guid.NewGuid();
        await using MqttSessionClient mqttExecutor = await ClientFactory.CreateSessionClientFromEnvAsync(executorId);

        await using CounterService counterService = new CounterService(applicationContext, mqttExecutor);
        await using MqttSessionClient mqttInvoker = await ClientFactory.CreateSessionClientFromEnvAsync();
        await using CounterClient counterClient = new CounterClient(applicationContext, mqttInvoker);

        await counterService.StartAsync(null, cancellationToken: CancellationToken.None);

        IncrementRequestPayload payload = new IncrementRequestPayload
        {
            IncrementValue = 1
        };

        var resp = await counterClient.ReadCounterAsync(executorId, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();
        Assert.Equal(0, resp.Response.CounterResponse);

        var resp2 = await counterClient.IncrementAsync(executorId, payload, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();
        Assert.Equal(1, resp2.Response.CounterResponse);

        var resp3 = await counterClient.ReadCounterAsync(executorId, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();
        Assert.Equal(1, resp3.Response.CounterResponse);

        await counterClient.IncrementAsync(executorId, payload).WithMetadata();
        await counterClient.IncrementAsync(executorId, payload).WithMetadata();
        await counterClient.IncrementAsync(executorId, payload).WithMetadata();

        var resp4 = await counterClient.ReadCounterAsync(executorId, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();
        Assert.Equal(4, resp4.Response.CounterResponse);

        await counterClient.ResetAsync(executorId).WithMetadata();

        var resp5 = await counterClient.ReadCounterAsync(executorId, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();
        Assert.Equal(0, resp5.Response.CounterResponse);

        await counterService.StopAsync();
        await counterClient.StopAsync();
    }
 
    [Fact]
    public async Task DuplicateCorrelationIDThrows()
    {
        ApplicationContext applicationContext = new ApplicationContext();
        string executorId = "counter-server-" + Guid.NewGuid();
        ApplicationContext appContext = new ApplicationContext();
        await using MqttSessionClient mqttExecutor = await ClientFactory.CreateSessionClientFromEnvAsync(executorId);

        await using CounterService counterService = new CounterService(applicationContext, mqttExecutor);
        await using MqttSessionClient mqttInvoker = await ClientFactory.CreateSessionClientFromEnvAsync();
        await using CounterClient counterClient = new CounterClient(appContext, mqttInvoker);

        await counterService.StartAsync(null, cancellationToken: CancellationToken.None);

        var resp = await counterClient.ReadCounterAsync(executorId, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();
        Assert.Equal(0, resp.Response.CounterResponse);
        
        CommandRequestMetadata reqMd2 = new();
        IncrementRequestPayload payload = new IncrementRequestPayload
        {
            IncrementValue = 1
        };
        Task[] tasks = new Task[2];
        for (int i = 0; i < tasks.Length; i++)
        {
            Task<ExtendedResponse<IncrementResponsePayload>> incrCounterTask = counterClient.IncrementAsync(executorId, payload, reqMd2).WithMetadata();
            tasks[i] = incrCounterTask;
        }
        var exception = await Assert.ThrowsAsync<AkriMqttException>(() => Task.WhenAll(tasks));
        Assert.Equal("Command 'increment' invocation failed due to duplicate request with same correlationId" ,exception.Message);

        await counterService.StopAsync();
        await counterClient.StopAsync();
    } 

    [Fact]
    public async Task WrongCorrelationIDThrows()
    {
        string executorId = "counter-server-" + Guid.NewGuid();
        ApplicationContext applicationContext = new ApplicationContext();
        await using MqttSessionClient mqttExecutor = await ClientFactory.CreateSessionClientFromEnvAsync(executorId);
        await using CounterService counterService = new CounterService(applicationContext, mqttExecutor);
        await counterService.StartAsync(null, cancellationToken: CancellationToken.None);

        await using MqttSessionClient mqttInvoker = await ClientFactory.CreateSessionClientFromEnvAsync("counter-client-bad");

        TaskCompletionSource<MqttApplicationMessageReceivedEventArgs> tcs = new();
        mqttInvoker.ApplicationMessageReceivedAsync += msg =>
        {
            tcs.SetResult(msg);
            return Task.CompletedTask;
        };

        var subscribeOptions = new MqttClientSubscribeOptions($"clients/counter-client-bad/rpc/command-samples/{executorId}/readCounter");
        await mqttInvoker.SubscribeAsync(subscribeOptions);

        var pubAck = await mqttInvoker.PublishAsync(
            new MqttApplicationMessage($"rpc/command-samples/{executorId}/readCounter")
            {
                CorrelationData = [0x01, 0xAA],
                ResponseTopic = $"clients/counter-client-bad/rpc/command-samples/{executorId}/readCounter",
                MessageExpiryInterval = 10,
            });

        MqttApplicationMessageReceivedEventArgs respMsg = await tcs.Task.WaitAsync(TimeSpan.FromMinutes(1));
        var userProps = respMsg.ApplicationMessage.UserProperties;
        Assert.Equal(7, userProps!.Count); // The user props are __stat, __stMsg, __protVer, __ts, $high_priority, __apErr, __propName.
        Assert.Equal("400", userProps.Where( p => p.Name == "__stat").First().Value);
        Assert.Equal("Correlation data bytes do not conform to a GUID.", userProps.FirstOrDefault(p => p.Name == "__stMsg")!.Value);
        Assert.Equal("Correlation Data", userProps.Where(p => p.Name == "__propName").First().Value);

        await counterService.StopAsync();
    }

    [Fact]
    public async Task CanReceiveApplicationErrorResponseInHeaders()
    {
        ApplicationContext applicationContext = new ApplicationContext();
        string executorId = "counter-server-" + Guid.NewGuid();
        await using MqttSessionClient mqttExecutor = await ClientFactory.CreateSessionClientFromEnvAsync(executorId);

        await using CounterService counterService = new CounterService(applicationContext, mqttExecutor);
        await using MqttSessionClient mqttInvoker = await ClientFactory.CreateSessionClientFromEnvAsync();
        await using CounterClient counterClient = new CounterClient(applicationContext, mqttInvoker);

        await counterService.StartAsync(null, cancellationToken: CancellationToken.None);

        int expectedNegativeValue = -1;
        IncrementRequestPayload payload = new IncrementRequestPayload
        {
            IncrementValue = expectedNegativeValue
        };

        var resp = await counterClient.IncrementAsync(executorId, payload, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();
        Assert.Equal(0, resp.Response.CounterResponse);
        Assert.NotNull(resp.ResponseMetadata);
        Assert.True(resp.TryGetApplicationError(out string? errorCode, out string? errorPayload));
        Assert.NotNull(errorCode);
        Assert.Equal(CounterService.NegativeValueArgumentErrorCode, errorCode);
        Assert.NotNull(errorPayload);
        CounterServiceApplicationError? deserializedErrorPayload = JsonSerializer.Deserialize<CounterServiceApplicationError>(errorPayload);
        Assert.NotNull(deserializedErrorPayload);
        Assert.Equal(expectedNegativeValue, deserializedErrorPayload.InvalidRequestArgumentValue);

        await counterService.StopAsync();
        await counterClient.StopAsync();
    }

    [Fact]
    public async Task TestRpcWithCloudEvents()
    {
        ApplicationContext applicationContext = new ApplicationContext();
        string executorId = "counter-server-" + Guid.NewGuid();
        await using MqttSessionClient mqttExecutor = await ClientFactory.CreateSessionClientFromEnvAsync(executorId);

        await using CounterService counterService = new CounterService(applicationContext, mqttExecutor);
        await using MqttSessionClient mqttInvoker = await ClientFactory.CreateSessionClientFromEnvAsync();
        await using CounterClient counterClient = new CounterClient(applicationContext, mqttInvoker);

        await counterService.StartAsync(null, cancellationToken: CancellationToken.None);

        IncrementRequestPayload payload = new IncrementRequestPayload
        {
            IncrementValue = 1
        };

        CloudEvent sentInvokeCloudEvent = new(new Uri("https://www.microsoft.com"), "someRpc.type")
        {
            DataSchema = "https://www.microsoft.com",
            Id = Guid.NewGuid().ToString(),
            Subject = "someSubject",
            Time = DateTime.UtcNow,
        };

        CommandRequestMetadata requestMetadata = new()
        {
            CloudEvent = sentInvokeCloudEvent,
        };

        var resp = await counterClient.ReadCounterAsync(executorId, requestMetadata, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();

        ExtendedCloudEvent? receivedResponseCloudEvent = resp.ResponseMetadata!.ExtendedCloudEvent;

        Assert.NotNull(counterService.ReceivedCloudEvent);
        Assert.NotNull(counterService.PublishedResponseCloudEvent);
        Assert.NotNull(receivedResponseCloudEvent);

        // Check that the cloud event sent by the invoker is read by the executor correctly
        Assert.Equal(sentInvokeCloudEvent.Id, counterService.ReceivedCloudEvent.Id);
        Assert.Equal(sentInvokeCloudEvent.Subject, counterService.ReceivedCloudEvent.Subject);
        Assert.Equal(sentInvokeCloudEvent.DataSchema, counterService.ReceivedCloudEvent.DataSchema);
        Assert.Equal(sentInvokeCloudEvent.Source, counterService.ReceivedCloudEvent.Source);
        Assert.Equal(sentInvokeCloudEvent.SpecVersion, counterService.ReceivedCloudEvent.SpecVersion);
        Assert.Equal(sentInvokeCloudEvent.Type, counterService.ReceivedCloudEvent.Type);
        Assert.Null(counterService.ReceivedCloudEvent.DataContentType);

        // Check that the cloud event sent by the executor in the response is read by the invoker correctly
        Assert.Equal(counterService.PublishedResponseCloudEvent.Id, receivedResponseCloudEvent.Id);
        Assert.Equal(counterService.PublishedResponseCloudEvent.Subject, receivedResponseCloudEvent.Subject);
        Assert.Equal(counterService.PublishedResponseCloudEvent.DataSchema, receivedResponseCloudEvent.DataSchema);
        Assert.Equal(counterService.PublishedResponseCloudEvent.Source, receivedResponseCloudEvent.Source);
        Assert.Equal(counterService.PublishedResponseCloudEvent.SpecVersion, receivedResponseCloudEvent.SpecVersion);
        Assert.Equal(counterService.PublishedResponseCloudEvent.Type, receivedResponseCloudEvent.Type);
        Assert.Equal("application/json", receivedResponseCloudEvent.DataContentType);
        
        sentInvokeCloudEvent.Subject = null;
        sentInvokeCloudEvent.Time = null;

        CommandRequestMetadata secondRequestMetadata = new()
        {
            CloudEvent = sentInvokeCloudEvent,
        };

        var resp2 = await counterClient.ReadCounterAsync(executorId, secondRequestMetadata, commandTimeout: TimeSpan.FromSeconds(30)).WithMetadata();

        ExtendedCloudEvent? receivedResponseCloudEvent2 = resp2.ResponseMetadata!.ExtendedCloudEvent;

        // Check that the cloud event sent by the invoker is read by the executor correctly
        Assert.NotNull(counterService.ReceivedCloudEvent);
        Assert.Null(counterService.ReceivedCloudEvent.Subject);
        Assert.Null(counterService.ReceivedCloudEvent.Time);

        // Check that the cloud event sent by the executor in the response is read by the invoker correctly
        Assert.NotNull(receivedResponseCloudEvent2);
        Assert.Null(receivedResponseCloudEvent2.Subject);
        Assert.Null(receivedResponseCloudEvent2.Time);

        await counterService.StopAsync();
        await counterClient.StopAsync();
    }

    [Fact]
    public async Task CanDisposeRpcClientsEvenIfMqttClientIsDisconnected()
    {
        ApplicationContext applicationContext = new ApplicationContext();
        string executorId = "counter-server-" + Guid.NewGuid();

        MqttSessionClientOptions clientOptions = new MqttSessionClientOptions()
        {
            ThrowIfUsedWhenSessionInactive = true
        };

        MqttSessionClient mqttExecutor = await ClientFactory.CreateSessionClientFromEnvAsync(executorId, clientOptions);
        CounterService counterService = new CounterService(applicationContext, mqttExecutor);
        MqttSessionClient mqttInvoker = await ClientFactory.CreateSessionClientFromEnvAsync("", clientOptions);
        CounterClient counterClient = new CounterClient(applicationContext, mqttInvoker);

        await counterService.StartAsync(null, cancellationToken: CancellationToken.None);

        await mqttExecutor.DisconnectAsync();
        await mqttInvoker.DisconnectAsync();

        await counterService.DisposeAsync();
        await counterClient.DisposeAsync();
    }

    [Fact]
    public async Task StopRpcClientsWithCancellationTokenThrows()
    {
        ApplicationContext applicationContext = new ApplicationContext();
        string executorId = "counter-server-" + Guid.NewGuid();

        MqttSessionClientOptions clientOptions = new MqttSessionClientOptions()
        {
            ThrowIfUsedWhenSessionInactive = false // Intentionally allow the pub/sub/unsub to hang after session client is closed
        };

        MqttSessionClient mqttExecutor = await ClientFactory.CreateSessionClientFromEnvAsync(executorId, clientOptions);
        CounterService counterService = new CounterService(applicationContext, mqttExecutor);
        MqttSessionClient mqttInvoker = await ClientFactory.CreateSessionClientFromEnvAsync("", clientOptions);
        CounterClient counterClient = new CounterClient(applicationContext, mqttInvoker);

        await counterService.StartAsync(null, cancellationToken: CancellationToken.None);

        await counterClient.StartAsync();
        var unused = await counterClient.ReadCounterAsync(executorId);

        await mqttExecutor.DisconnectAsync();
        await mqttInvoker.DisconnectAsync();

        using CancellationTokenSource cts1 = new CancellationTokenSource(10);
        try
        {
            await counterService.StopAsync(cts1.Token);
            Assert.Fail("Expected the operation to throw an OperationCanceledException or TaskCanceledException, but no exception was thrown");
        }
        catch (OperationCanceledException)
        {
            // Expected exception was thrown, continue
        }

        using CancellationTokenSource cts2 = new CancellationTokenSource(10);
        try
        {
            await counterClient.StopAsync(cts2.Token);
            Assert.Fail("Expected the operation to throw an OperationCanceledException or TaskCanceledException, but no exception was thrown");
        }
        catch (OperationCanceledException)
        {
            // Expected exception was thrown, continue
        }

        await counterService.DisposeAsync(true, CancellationToken.None);
        await counterClient.DisposeAsync(true, CancellationToken.None);
    }

    [Fact]
    public async Task CanDisposeRpcClientsWithCancellationToken()
    {
        ApplicationContext applicationContext = new ApplicationContext();
        string executorId = "counter-server-" + Guid.NewGuid();

        MqttSessionClientOptions clientOptions = new MqttSessionClientOptions()
        {
            ThrowIfUsedWhenSessionInactive = false // Intentionally allow the pub/sub/unsub to hang after session client is closed
        };

        MqttSessionClient mqttExecutor = await ClientFactory.CreateSessionClientFromEnvAsync(executorId, clientOptions);
        CounterService counterService = new CounterService(applicationContext, mqttExecutor);
        MqttSessionClient mqttInvoker = await ClientFactory.CreateSessionClientFromEnvAsync("", clientOptions);
        CounterClient counterClient = new CounterClient(applicationContext, mqttInvoker);

        await counterService.StartAsync(null, cancellationToken: CancellationToken.None);

        await mqttExecutor.DisconnectAsync();
        await mqttInvoker.DisconnectAsync();

        await counterService.DisposeAsync(true, CancellationToken.None); //TODO optional param in codegen layer
        await counterClient.DisposeAsync(true, CancellationToken.None);
    }
}
