namespace NATS.Server.Raft;

public interface IRaftTransport
{
    Task<IReadOnlyList<AppendResult>> AppendEntriesAsync(string leaderId, IReadOnlyList<string> followerIds, RaftLogEntry entry, CancellationToken ct);
    Task<VoteResponse> RequestVoteAsync(string candidateId, string voterId, VoteRequest request, CancellationToken ct);
    Task InstallSnapshotAsync(string leaderId, string followerId, RaftSnapshot snapshot, CancellationToken ct);

    /// <summary>
    /// Sends a TimeoutNow RPC to the target follower, asking it to immediately start
    /// an election and bypass its election timer. Used for leadership transfer.
    /// Go reference: raft.go sendTimeoutNow
    /// </summary>
    Task SendTimeoutNowAsync(string leaderId, string targetId, ulong term, CancellationToken ct);

    /// <summary>
    /// Sends heartbeat RPCs to all listed followers and records their acknowledgement by
    /// updating the leader's peer <see cref="RaftPeerState.LastContact"/> timestamps.
    /// Used by <see cref="RaftNode.ReadIndexAsync"/> to confirm leadership quorum before
    /// serving a linearizable read.
    /// Go reference: raft.go — leader sends AppendEntries (empty) to confirm quorum for reads.
    /// </summary>
    Task SendHeartbeatAsync(string leaderId, IReadOnlyList<string> followerIds, int term, Action<string> onAck, CancellationToken ct);
}

public sealed class InMemoryRaftTransport : IRaftTransport
{
    private readonly Dictionary<string, RaftNode> _nodes = new(StringComparer.Ordinal);

    public void Register(RaftNode node)
    {
        _nodes[node.Id] = node;
    }

    public Task<IReadOnlyList<AppendResult>> AppendEntriesAsync(string leaderId, IReadOnlyList<string> followerIds, RaftLogEntry entry, CancellationToken ct)
    {
        var results = new List<AppendResult>(followerIds.Count);
        foreach (var followerId in followerIds)
        {
            if (_nodes.TryGetValue(followerId, out var node))
            {
                node.ReceiveReplicatedEntry(entry);
                results.Add(new AppendResult { FollowerId = followerId, Success = true });
            }
            else
            {
                results.Add(new AppendResult { FollowerId = followerId, Success = false });
            }
        }

        return Task.FromResult<IReadOnlyList<AppendResult>>(results);
    }

    public Task<VoteResponse> RequestVoteAsync(string candidateId, string voterId, VoteRequest request, CancellationToken ct)
    {
        if (_nodes.TryGetValue(voterId, out var node))
            return Task.FromResult(node.GrantVote(request.Term, string.IsNullOrWhiteSpace(request.CandidateId) ? candidateId : request.CandidateId));

        return Task.FromResult(new VoteResponse { Granted = false });
    }

    public async Task InstallSnapshotAsync(string leaderId, string followerId, RaftSnapshot snapshot, CancellationToken ct)
    {
        _ = leaderId;
        if (_nodes.TryGetValue(followerId, out var node))
            await node.InstallSnapshotAsync(snapshot, ct);
    }

    public async Task AppendHeartbeatAsync(string leaderId, IReadOnlyList<string> followerIds, int term, CancellationToken ct)
    {
        _ = leaderId;
        foreach (var followerId in followerIds)
        {
            if (_nodes.TryGetValue(followerId, out var node))
                node.ReceiveHeartbeat(term);
        }

        await Task.CompletedTask;
    }

    /// <summary>
    /// Delivers a heartbeat to each listed follower (simulating an AppendEntries RPC with no
    /// entries), invokes <paramref name="onAck"/> for each reachable follower so the leader can
    /// update its peer <see cref="RaftPeerState.LastContact"/> timestamps, then returns.
    /// Unreachable followers (not registered in the transport) produce no acknowledgement.
    /// Go reference: raft.go — leader sends empty AppendEntries to confirm quorum for reads.
    /// </summary>
    public Task SendHeartbeatAsync(string leaderId, IReadOnlyList<string> followerIds, int term, Action<string> onAck, CancellationToken ct)
    {
        _ = leaderId;
        foreach (var followerId in followerIds)
        {
            if (_nodes.TryGetValue(followerId, out var node))
            {
                node.ReceiveHeartbeat(term);
                onAck(followerId);
            }
        }

        return Task.CompletedTask;
    }

    /// <summary>
    /// Delivers a TimeoutNow RPC to the target node, causing it to start an election
    /// immediately by calling <see cref="RaftNode.ReceiveTimeoutNow"/>.
    /// If the target is not registered (simulating an unreachable peer), does nothing.
    /// Go reference: raft.go sendTimeoutNow / processTimeoutNow
    /// </summary>
    public Task SendTimeoutNowAsync(string leaderId, string targetId, ulong term, CancellationToken ct)
    {
        _ = leaderId;
        if (_nodes.TryGetValue(targetId, out var node))
            node.ReceiveTimeoutNow(term);
        return Task.CompletedTask;
    }
}
