using NATS.Server.Raft;

namespace NATS.Server.JetStream.Cluster;

public sealed class StreamReplicaGroup
{
    private readonly List<RaftNode> _nodes;

    public string StreamName { get; }
    public IReadOnlyList<RaftNode> Nodes => _nodes;
    public RaftNode Leader { get; private set; }

    public StreamReplicaGroup(string streamName, int replicas)
    {
        StreamName = streamName;

        var nodeCount = Math.Max(replicas, 1);
        _nodes = Enumerable.Range(1, nodeCount)
            .Select(i => new RaftNode($"{streamName.ToLowerInvariant()}-r{i}"))
            .ToList();

        foreach (var node in _nodes)
            node.ConfigureCluster(_nodes);

        Leader = ElectLeader(_nodes[0]);
    }

    public async ValueTask<long> ProposeAsync(string command, CancellationToken ct)
    {
        if (!Leader.IsLeader)
            Leader = ElectLeader(SelectNextCandidate(Leader));

        return await Leader.ProposeAsync(command, ct);
    }

    public Task StepDownAsync(CancellationToken ct)
    {
        _ = ct;
        var previous = Leader;
        previous.RequestStepDown();
        Leader = ElectLeader(SelectNextCandidate(previous));
        return Task.CompletedTask;
    }

    private RaftNode SelectNextCandidate(RaftNode currentLeader)
    {
        if (_nodes.Count == 1)
            return _nodes[0];

        var index = _nodes.FindIndex(n => n.Id == currentLeader.Id);
        if (index < 0)
            return _nodes[0];

        return _nodes[(index + 1) % _nodes.Count];
    }

    private RaftNode ElectLeader(RaftNode candidate)
    {
        candidate.StartElection(_nodes.Count);
        foreach (var voter in _nodes.Where(n => n.Id != candidate.Id))
            candidate.ReceiveVote(voter.GrantVote(candidate.Term), _nodes.Count);

        return candidate;
    }
}