using NATS.Server.Raft;

namespace NATS.Server.Tests;

public class RaftElectionTests
{
    [Fact]
    public async Task Candidate_becomes_leader_after_majority_votes()
    {
        var cluster = RaftTestCluster.Create(3);
        var leader = await cluster.ElectLeaderAsync();

        leader.Role.ShouldBe(RaftRole.Leader);
        leader.Term.ShouldBe(1);
    }
}

internal sealed class RaftTestCluster
{
    public List<RaftNode> Nodes { get; }
    public RaftNode Leader { get; private set; }
    public RaftNode LaggingFollower { get; private set; }

    private RaftTestCluster(List<RaftNode> nodes)
    {
        Nodes = nodes;
        Leader = nodes[0];
        LaggingFollower = nodes[^1];
    }

    public static RaftTestCluster Create(int nodes)
    {
        var created = Enumerable.Range(1, nodes).Select(i => new RaftNode($"n{i}")).ToList();
        foreach (var node in created)
            node.ConfigureCluster(created);
        return new RaftTestCluster(created);
    }

    public Task<RaftNode> ElectLeaderAsync()
    {
        var candidate = Nodes[0];
        candidate.StartElection(Nodes.Count);

        foreach (var voter in Nodes.Skip(1))
            candidate.ReceiveVote(voter.GrantVote(candidate.Term));

        Leader = candidate;
        return Task.FromResult(candidate);
    }

    public async Task WaitForAppliedAsync(long index)
    {
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(2));
        while (!timeout.IsCancellationRequested)
        {
            if (Nodes.All(n => n.AppliedIndex >= index))
                return;

            await Task.Delay(20, timeout.Token).ContinueWith(_ => { }, TaskScheduler.Default);
        }
    }

    public async Task GenerateCommittedEntriesAsync(int count)
    {
        var leader = await ElectLeaderAsync();
        for (int i = 0; i < count; i++)
            _ = await leader.ProposeAsync($"cmd-{i}", default);
    }

    public Task RestartLaggingFollowerAsync()
    {
        LaggingFollower = Nodes[^1];
        LaggingFollower.AppliedIndex = 0;
        return Task.CompletedTask;
    }

    public async Task WaitForFollowerCatchupAsync()
    {
        var snapshot = await Leader.CreateSnapshotAsync(default);
        await LaggingFollower.InstallSnapshotAsync(snapshot, default);
    }
}