using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Threading.Tasks;

using Microsoft.Extensions.Logging;
using Microsoft.VisualStudio.TestTools.UnitTesting;

using Tgstation.Server.Common.Extensions;
using Tgstation.Server.Host.Configuration;
using Tgstation.Server.Host.Core;
using Tgstation.Server.Host.IO;

namespace Tgstation.Server.Host.Swarm.Tests
{
	[TestClass]
	public sealed class TestSwarmProtocol
	{
		HashSet<ushort> usedPorts = new();
		ILoggerFactory loggerFactory;
		ILogger logger;

		ISeekableFileStreamProvider updateFileStreamProvider;

		[TestInitialize]
		public async Task Initialize()
		{
			loggerFactory = LoggerFactory.Create(builder =>
			{
				builder.SetMinimumLevel(LogLevel.Trace);
				builder.AddConsole();
			});

			logger = loggerFactory.CreateLogger<TestSwarmProtocol>();

			var ms = new MemoryStream(new byte[] { 1, 2, 3, 4 });

			updateFileStreamProvider = new BufferedFileStreamProvider(ms);
			await updateFileStreamProvider.GetResult(default);
		}

		[TestCleanup]
		public async Task Shutdown()
		{
			usedPorts.Clear();
			loggerFactory.Dispose();
			await updateFileStreamProvider.DisposeAsync();
		}

		[TestMethod]
		public async Task TestInitHappensInstantlyWhenControllerIsInitialized()
		{
			await using var controller = GenNode();
			await using var node = GenNode(controller);

			TestableSwarmNode.Link(controller, node);

			controller.RpcMapper.AsyncRequests = false;
			node.RpcMapper.AsyncRequests = false;

			var controllerInit = controller.TryInit();
			Assert.IsTrue(controllerInit.IsCompleted);
			Assert.AreEqual(SwarmRegistrationResult.Success, await controllerInit);

			var nodeInit = node.TryInit();
			Assert.IsTrue(nodeInit.IsCompleted);
			Assert.AreEqual(SwarmRegistrationResult.Success, await nodeInit);
		}

		[TestMethod]
		public async Task TestNodeInitializeDoesNotWorkWithoutController()
		{
			await using var controller = GenNode();
			await using var node1 = GenNode(controller);
			await using var node2 = GenNode(controller);

			TestableSwarmNode.Link(controller, node1, node2);

			Assert.AreEqual(SwarmRegistrationResult.CommunicationFailure, await node1.TryInit());

			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());

			Assert.AreEqual(SwarmRegistrationResult.Success, await node2.TryInit());
		}

		[TestMethod]
		public async Task TestUpdateCantProceedWithoutAllNodes()
		{
			await using var controller = GenNode();
			await using var node1 = GenNode(controller);
			await using var node2 = GenNode(controller);

			TestableSwarmNode.Link(controller, node1, node2);

			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
			Assert.AreEqual(SwarmRegistrationResult.Success, await node2.TryInit());

			await DelayMax(() =>
			{
				Assert.AreEqual(2, controller.Service.GetSwarmServers().Count);
				Assert.AreEqual(2, node2.Service.GetSwarmServers().Count);
			}, 5);

			Assert.AreEqual(SwarmPrepareResult.Failure, await controller.Service.PrepareUpdate(updateFileStreamProvider, new Version(4, 3, 2), default));

			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());

			await DelayMax(() =>
			{
				Assert.AreEqual(3, controller.Service.GetSwarmServers().Count);
				Assert.AreEqual(3, node2.Service.GetSwarmServers().Count);
				Assert.AreEqual(3, node1.Service.GetSwarmServers().Count);
			}, 5);

			Assert.AreEqual(SwarmPrepareResult.SuccessHoldProviderUntilCommit, await controller.Service.PrepareUpdate(updateFileStreamProvider, new Version(4, 3, 2), default));
		}

		[TestMethod]
		public async Task TestServersReconnectAfterReboot()
		{
			await using var controller = GenNode();
			await using var node1 = GenNode(controller);

			TestableSwarmNode.Link(controller, node1);
			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());

			node1.RpcMapper.AsyncRequests = false;
			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());
			node1.RpcMapper.AsyncRequests = true;

			Assert.AreEqual(2, controller.Service.GetSwarmServers().Count);
			Assert.AreEqual(1, node1.Service.GetSwarmServers().Count);

			await DelayMax(() => Assert.AreEqual(2, node1.Service.GetSwarmServers().Count));

			await controller.SimulateReboot(default);
			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
			Assert.AreEqual(1, controller.Service.GetSwarmServers().Count);

			// Remember node's async delayer fires every millisecond
			// server list is updated in health check loop
			await DelayMax(() =>
			{
				Assert.AreEqual(2, controller.Service.GetSwarmServers().Count);
				Assert.AreEqual(2, node1.Service.GetSwarmServers().Count);
			});

			await node1.SimulateReboot(default);
			// health check timeout
			await DelayMax(() => Assert.AreEqual(1, controller.Service.GetSwarmServers().Count));

			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());

			Assert.AreEqual(2, controller.Service.GetSwarmServers().Count);
			await DelayMax(() => Assert.AreEqual(2, node1.Service.GetSwarmServers().Count));
		}

		[TestMethod]
		public async Task TestCommitFailsIfNotPrepared()
		{
			await using var controller = GenNode();
			await using var node1 = GenNode(controller);

			TestableSwarmNode.Link(controller, node1);

			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());

			Assert.AreEqual(SwarmCommitResult.AbortUpdate, await node1.Service.CommitUpdate(default));
		}

		[TestMethod]
		public async Task TestPrepareDifferentVersionsFails()
		{
			await using var controller = GenNode();
			await using var node1 = GenNode(controller);

			TestableSwarmNode.Link(controller, node1);

			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());

			await DelayMax(() =>
			{
				Assert.AreEqual(2, controller.Service.GetSwarmServers().Count);
				Assert.AreEqual(2, node1.Service.GetSwarmServers().Count);
			}, 5);

			Assert.AreEqual(SwarmPrepareResult.SuccessHoldProviderUntilCommit, await node1.Service.PrepareUpdate(updateFileStreamProvider, new Version(4, 3, 2), default));
			Assert.AreEqual(SwarmPrepareResult.Failure, await node1.Service.PrepareUpdate(updateFileStreamProvider, new Version(4, 3, 3), default));
		}

		[TestMethod]
		public async Task TestSimultaneousPrepareDifferentVersionsFailsControllerFirst()
		{
			await TestSimultaneousPrepareDifferentVersionsFails(true);
		}

		[TestMethod]
		public async Task TestSimultaneousPrepareDifferentVersionsFailsNodeFirst()
		{
			await TestSimultaneousPrepareDifferentVersionsFails(false);
		}

		async ValueTask TestSimultaneousPrepareDifferentVersionsFails(bool prepControllerFirst)
		{
			await using var controller = GenNode();
			await using var node1 = GenNode(controller);

			TestableSwarmNode.Link(controller, node1);

			controller.UpdateResult = ServerUpdateResult.UpdateInProgress;
			node1.UpdateResult = ServerUpdateResult.UpdateInProgress;

			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());

			ValueTask<SwarmPrepareResult> controllerPrepareTask = ValueTask.FromResult(SwarmPrepareResult.SuccessProviderNotRequired), nodePrepareTask;
			if (prepControllerFirst)
				controllerPrepareTask = controller.Service.PrepareUpdate(updateFileStreamProvider, new Version(4, 3, 2), default);

			nodePrepareTask = node1.Service.PrepareUpdate(updateFileStreamProvider, new Version(4, 3, 3), default);

			if (!prepControllerFirst)
				controllerPrepareTask = controller.Service.PrepareUpdate(updateFileStreamProvider, new Version(4, 3, 2), default);

			await Task.Yield();

			await ValueTaskExtensions.WhenAll(controllerPrepareTask, nodePrepareTask);

			Task<SwarmCommitResult> nodeCommitTask = null, controllerCommitTask = null;

			var controllerPrepped = await controllerPrepareTask;
			var nodePrepped = await nodePrepareTask;
			if (controllerPrepped == SwarmPrepareResult.Failure && nodePrepped == SwarmPrepareResult.Failure)
				return; // all is good with the world

			// We have to be fair to the system here...
			// Requests could still be in flight, things could get aborted
			// Each server STILL has to download the update package which takes time
			// Give it 1 second to get its shit together
			await Task.Delay(TimeSpan.FromSeconds(1));

			if (controllerPrepped != SwarmPrepareResult.Failure)
			{
				Assert.AreEqual(SwarmPrepareResult.SuccessHoldProviderUntilCommit, controllerPrepped);
				controllerCommitTask = controller.Service.CommitUpdate(default).AsTask();
			}

			if (nodePrepped != SwarmPrepareResult.Failure)
			{
				Assert.AreEqual(SwarmPrepareResult.SuccessHoldProviderUntilCommit, nodePrepped);
				nodeCommitTask = node1.Service.CommitUpdate(default).AsTask();
			}

			await Task.Yield();

			await Task.WhenAll(controllerCommitTask ?? Task.CompletedTask, nodeCommitTask ?? Task.CompletedTask);

			if (controllerPrepped != SwarmPrepareResult.Failure)
				Assert.AreEqual(SwarmCommitResult.AbortUpdate, await controllerCommitTask);

			if (nodePrepped != SwarmPrepareResult.Failure)
				Assert.AreEqual(SwarmCommitResult.AbortUpdate, await nodeCommitTask);
		}

		TestableSwarmNode GenNode(TestableSwarmNode controller = null, Version version = null)
		{
			ushort randPort;
			do
			{
				randPort = (ushort)(Random.Shared.Next() % UInt16.MaxValue);
			}
			while (randPort == 0 || !usedPorts.Add(randPort));

			var config = new SwarmConfiguration
			{
				Address = new Uri($"http://127.0.0.1:{randPort}"),
				ControllerAddress = controller?.Config.Address,
				Identifier = $"{(controller == null ? "Controller" : "Node")}{usedPorts.Count}",
				PrivateKey = "asdf",
				UpdateRequiredNodeCount = (uint)usedPorts.Count - 1,
			};

			return new TestableSwarmNode(loggerFactory, config, version);
		}

		async Task DelayMax(Action assertion, long seconds = 1)
		{
			var id = Guid.NewGuid();
			logger.LogInformation("Begin DelayMax {id}: {seconds}", id, seconds);
			var stopwatch = Stopwatch.StartNew();
			do
			{
				try
				{
					assertion();
					logger.LogInformation("End DelayMax {id} after {milliseconds}ms", id, stopwatch.ElapsedMilliseconds);
					return;
				}
				catch (AssertFailedException)
				{
				}

				await Task.Delay(TimeSpan.FromMilliseconds(100));
			}
			while (stopwatch.ElapsedMilliseconds < (seconds * 1000));

			try
			{
				assertion();
			}
			catch
			{
				logger.LogError("Fail DelayMax {id} after {milliseconds}ms", id, stopwatch.ElapsedMilliseconds);
				throw;
			}

			logger.LogInformation("End DelayMax {id} after {milliseconds}ms", id, stopwatch.ElapsedMilliseconds);
		}
	}
}