using NATS.Server.Raft;

namespace NATS.Server.Raft.Tests.Raft;

/// <summary>
/// Election behavior tests covering leader election, vote mechanics, term handling,
/// candidate stepdown, split vote scenarios, and network partition leader stepdown.
/// Go: TestNRGSimple, TestNRGSimpleElection, TestNRGInlineStepdown,
///     TestNRGRecoverFromFollowingNoLeader, TestNRGStepDownOnSameTermDoesntClearVote,
///     TestNRGAssumeHighTermAfterCandidateIsolation in server/raft_test.go.
/// </summary>
public class RaftElectionTests
{
    // -- Helpers (self-contained, no shared TestHelpers class) --

    private static (RaftNode[] nodes, InMemoryRaftTransport transport) CreateCluster(int size)
    {
        var transport = new InMemoryRaftTransport();
        var nodes = Enumerable.Range(1, size)
            .Select(i => new RaftNode($"n{i}", transport))
            .ToArray();
        foreach (var node in nodes)
        {
            transport.Register(node);
            node.ConfigureCluster(nodes);
        }
        return (nodes, transport);
    }

    private static RaftNode ElectLeader(RaftNode[] nodes)
    {
        var candidate = nodes[0];
        candidate.StartElection(nodes.Length);
        foreach (var voter in nodes.Skip(1))
            candidate.ReceiveVote(voter.GrantVote(candidate.Term, candidate.Id), nodes.Length);
        return candidate;
    }

    // Go: TestNRGSimple server/raft_test.go:35
    [Fact]
    public void Single_node_becomes_leader_automatically()
    {
        var node = new RaftNode("solo");
        node.StartElection(clusterSize: 1);

        node.IsLeader.ShouldBeTrue();
        node.Role.ShouldBe(RaftRole.Leader);
        node.Term.ShouldBe(1);
    }

    // Go: TestNRGSimple server/raft_test.go:35
    [Fact]
    public void Three_node_cluster_elects_leader()
    {
        var (nodes, _) = CreateCluster(3);
        var leader = ElectLeader(nodes);

        leader.IsLeader.ShouldBeTrue();
        leader.Role.ShouldBe(RaftRole.Leader);
        nodes.Count(n => n.IsLeader).ShouldBe(1);
        nodes.Count(n => !n.IsLeader).ShouldBe(2);
    }

    // Go: TestNRGSimpleElection server/raft_test.go:296
    [Fact]
    public void Five_node_cluster_elects_leader_with_quorum()
    {
        var (nodes, _) = CreateCluster(5);
        var leader = ElectLeader(nodes);

        leader.IsLeader.ShouldBeTrue();
        nodes.Count(n => n.IsLeader).ShouldBe(1);
        nodes.Count(n => !n.IsLeader).ShouldBe(4);
    }

    // Go: TestNRGSimpleElection server/raft_test.go:296
    [Fact]
    public void Election_increments_term()
    {
        var (nodes, _) = CreateCluster(3);
        var candidate = nodes[0];

        candidate.Term.ShouldBe(0);
        candidate.StartElection(nodes.Length);
        candidate.Term.ShouldBe(1);
    }

    // Go: TestNRGSimpleElection server/raft_test.go:296
    [Fact]
    public void Candidate_votes_for_self_on_election_start()
    {
        var node = new RaftNode("n1");
        node.StartElection(clusterSize: 3);

        node.Role.ShouldBe(RaftRole.Candidate);
        node.TermState.VotedFor.ShouldBe("n1");
    }

    // Go: TestNRGSimpleElection server/raft_test.go:296
    [Fact]
    public void Candidate_needs_majority_to_become_leader()
    {
        var (nodes, _) = CreateCluster(3);
        var candidate = nodes[0];

        candidate.StartElection(nodes.Length);
        // Only self-vote, not enough for majority in 3-node cluster
        candidate.IsLeader.ShouldBeFalse();
        candidate.Role.ShouldBe(RaftRole.Candidate);

        // One more vote gives majority (2 out of 3)
        var vote = nodes[1].GrantVote(candidate.Term, candidate.Id);
        vote.Granted.ShouldBeTrue();
        candidate.ReceiveVote(vote, nodes.Length);
        candidate.IsLeader.ShouldBeTrue();
    }

    // Go: TestNRGSimpleElection server/raft_test.go:296
    [Fact]
    public void Denied_vote_does_not_advance_to_leader()
    {
        var node = new RaftNode("n1");
        node.StartElection(clusterSize: 5);
        node.IsLeader.ShouldBeFalse();

        // Receive denied votes
        node.ReceiveVote(new VoteResponse { Granted = false }, clusterSize: 5);
        node.ReceiveVote(new VoteResponse { Granted = false }, clusterSize: 5);
        node.IsLeader.ShouldBeFalse();
    }

    // Go: TestNRGSimpleElection server/raft_test.go:296
    [Fact]
    public void Vote_granted_for_same_term_and_candidate()
    {
        var voter = new RaftNode("voter");
        var response = voter.GrantVote(term: 1, candidateId: "candidate-a");
        response.Granted.ShouldBeTrue();
        voter.TermState.VotedFor.ShouldBe("candidate-a");
    }

    // Go: TestNRGStepDownOnSameTermDoesntClearVote server/raft_test.go:447
    [Fact]
    public void Vote_denied_for_same_term_different_candidate()
    {
        var voter = new RaftNode("voter");
        // Vote for candidate-a in term 1
        voter.GrantVote(term: 1, candidateId: "candidate-a").Granted.ShouldBeTrue();

        // Attempt to vote for candidate-b in same term should fail
        var response = voter.GrantVote(term: 1, candidateId: "candidate-b");
        response.Granted.ShouldBeFalse();
        voter.TermState.VotedFor.ShouldBe("candidate-a");
    }

    // Go: processVoteRequest in server/raft.go — stale term rejection
    [Fact]
    public void Vote_denied_for_stale_term()
    {
        var voter = new RaftNode("voter");
        voter.TermState.CurrentTerm = 5;

        var response = voter.GrantVote(term: 3, candidateId: "candidate");
        response.Granted.ShouldBeFalse();
    }

    // Go: processVoteRequest in server/raft.go — higher term resets vote
    [Fact]
    public void Vote_granted_for_higher_term_resets_previous_vote()
    {
        var voter = new RaftNode("voter");
        voter.GrantVote(term: 1, candidateId: "candidate-a").Granted.ShouldBeTrue();
        voter.TermState.VotedFor.ShouldBe("candidate-a");

        // Higher term should clear previous vote and grant new one
        var response = voter.GrantVote(term: 2, candidateId: "candidate-b");
        response.Granted.ShouldBeTrue();
        voter.TermState.VotedFor.ShouldBe("candidate-b");
        voter.TermState.CurrentTerm.ShouldBe(2);
    }

    // Go: TestNRGInlineStepdown server/raft_test.go:194
    [Fact]
    public void Leader_stepdown_transitions_to_follower()
    {
        var node = new RaftNode("n1");
        node.StartElection(clusterSize: 1);
        node.IsLeader.ShouldBeTrue();

        node.RequestStepDown();
        node.IsLeader.ShouldBeFalse();
        node.Role.ShouldBe(RaftRole.Follower);
    }

    // Go: TestNRGInlineStepdown server/raft_test.go:194
    [Fact]
    public void Stepdown_clears_votes_received()
    {
        var (nodes, _) = CreateCluster(3);
        var leader = ElectLeader(nodes);
        leader.IsLeader.ShouldBeTrue();

        leader.RequestStepDown();
        leader.Role.ShouldBe(RaftRole.Follower);
        leader.TermState.VotedFor.ShouldBeNull();
    }

    // Go: TestNRGRecoverFromFollowingNoLeader server/raft_test.go:154
    [Fact]
    public void Candidate_stepdown_on_higher_term_heartbeat()
    {
        var node = new RaftNode("n1");
        node.StartElection(clusterSize: 3);
        node.Role.ShouldBe(RaftRole.Candidate);
        node.Term.ShouldBe(1);

        // Receive heartbeat with higher term
        node.ReceiveHeartbeat(term: 5);
        node.Role.ShouldBe(RaftRole.Follower);
        node.Term.ShouldBe(5);
    }

    // Go: TestNRGRecoverFromFollowingNoLeader server/raft_test.go:154
    [Fact]
    public void Leader_stepdown_on_higher_term_heartbeat()
    {
        var node = new RaftNode("n1");
        node.StartElection(clusterSize: 1);
        node.IsLeader.ShouldBeTrue();
        node.Term.ShouldBe(1);

        node.ReceiveHeartbeat(term: 10);
        node.Role.ShouldBe(RaftRole.Follower);
        node.Term.ShouldBe(10);
    }

    // Go: TestNRGRecoverFromFollowingNoLeader server/raft_test.go:154
    [Fact]
    public void Heartbeat_with_lower_term_ignored()
    {
        var node = new RaftNode("n1");
        node.StartElection(clusterSize: 1);
        node.IsLeader.ShouldBeTrue();
        node.Term.ShouldBe(1);

        node.ReceiveHeartbeat(term: 0);
        node.IsLeader.ShouldBeTrue();
        node.Term.ShouldBe(1);
    }

    // Go: TestNRGAssumeHighTermAfterCandidateIsolation server/raft_test.go:662
    [Fact]
    public void Split_vote_forces_reelection_with_higher_term()
    {
        var (nodes, _) = CreateCluster(3);

        // First election: n1 starts but only gets self-vote
        nodes[0].StartElection(nodes.Length);
        nodes[0].Role.ShouldBe(RaftRole.Candidate);
        nodes[0].Term.ShouldBe(1);

        // n2 also starts election concurrently (split vote scenario)
        nodes[1].StartElection(nodes.Length);
        nodes[1].Role.ShouldBe(RaftRole.Candidate);
        nodes[1].Term.ShouldBe(1);

        // Neither gets majority, so no leader
        nodes.Count(n => n.IsLeader).ShouldBe(0);

        // n1 starts new election in higher term
        nodes[0].StartElection(nodes.Length);
        nodes[0].Term.ShouldBe(2);

        // Now n2 and n3 grant votes
        var v2 = nodes[1].GrantVote(nodes[0].Term, nodes[0].Id);
        v2.Granted.ShouldBeTrue();
        nodes[0].ReceiveVote(v2, nodes.Length);
        nodes[0].IsLeader.ShouldBeTrue();
    }

    // Go: TestNRGAssumeHighTermAfterCandidateIsolation server/raft_test.go:662
    [Fact]
    public void Isolated_candidate_with_high_term_forces_term_update()
    {
        var (nodes, transport) = CreateCluster(3);
        var leader = ElectLeader(nodes);
        leader.IsLeader.ShouldBeTrue();

        // Simulate follower isolation: bump its term high
        var follower = nodes.First(n => !n.IsLeader);
        follower.TermState.CurrentTerm = 100;

        // When the isolated node's vote request reaches others,
        // they should update their term even if they don't grant the vote
        var voteReq = new VoteRequest { Term = 100, CandidateId = follower.Id };

        foreach (var node in nodes.Where(n => n.Id != follower.Id))
        {
            var resp = node.GrantVote(voteReq.Term, voteReq.CandidateId);
            // Term should update to 100 regardless of vote grant
            node.TermState.CurrentTerm.ShouldBe(100);
        }
    }

    // Go: TestNRGRecoverFromFollowingNoLeader server/raft_test.go:154
    [Fact]
    public void Re_election_after_leader_stepdown()
    {
        var (nodes, _) = CreateCluster(3);
        var leader = ElectLeader(nodes);
        leader.IsLeader.ShouldBeTrue();
        leader.Term.ShouldBe(1);

        // Leader steps down
        leader.RequestStepDown();
        leader.IsLeader.ShouldBeFalse();

        // New election with a different candidate — term increments from current
        var newCandidate = nodes.First(n => n.Id != leader.Id);
        newCandidate.StartElection(nodes.Length);
        newCandidate.Term.ShouldBe(2); // was 1 from first election, incremented to 2

        foreach (var voter in nodes.Where(n => n.Id != newCandidate.Id))
        {
            var vote = voter.GrantVote(newCandidate.Term, newCandidate.Id);
            newCandidate.ReceiveVote(vote, nodes.Length);
        }

        newCandidate.IsLeader.ShouldBeTrue();
    }

    // Go: TestNRGHeartbeatOnLeaderChange server/raft_test.go:708
    [Fact]
    public void Multiple_sequential_elections_increment_term()
    {
        var node = new RaftNode("n1");

        node.StartElection(clusterSize: 1);
        node.Term.ShouldBe(1);

        node.RequestStepDown();
        node.StartElection(clusterSize: 1);
        node.Term.ShouldBe(2);

        node.RequestStepDown();
        node.StartElection(clusterSize: 1);
        node.Term.ShouldBe(3);
    }

    // Go: TestNRGSimpleElection server/raft_test.go:296 — transport-based vote request
    [Fact]
    public async Task Transport_based_vote_request()
    {
        var (nodes, transport) = CreateCluster(3);

        var candidate = nodes[0];
        candidate.StartElection(nodes.Length);

        // Use transport to request votes
        var voteReq = new VoteRequest { Term = candidate.Term, CandidateId = candidate.Id };

        foreach (var voter in nodes.Skip(1))
        {
            var resp = await transport.RequestVoteAsync(candidate.Id, voter.Id, voteReq, default);
            candidate.ReceiveVote(resp, nodes.Length);
        }

        candidate.IsLeader.ShouldBeTrue();
    }

    // Go: TestNRGCandidateDoesntRevertTermAfterOldAE server/raft_test.go:792
    [Fact]
    public void Candidate_does_not_revert_term_on_stale_heartbeat()
    {
        var node = new RaftNode("n1");
        node.StartElection(clusterSize: 3);
        node.Term.ShouldBe(1);

        // Start another election to bump term
        node.StartElection(clusterSize: 3);
        node.Term.ShouldBe(2);

        // Receiving heartbeat from older term should not revert
        node.ReceiveHeartbeat(term: 1);
        node.Term.ShouldBe(2);
    }

    // Go: TestNRGCandidateDontStepdownDueToLeaderOfPreviousTerm server/raft_test.go:972
    [Fact]
    public void Candidate_does_not_stepdown_from_old_term_heartbeat()
    {
        var node = new RaftNode("n1");
        node.TermState.CurrentTerm = 10;
        node.StartElection(clusterSize: 3);
        node.Term.ShouldBe(11);
        node.Role.ShouldBe(RaftRole.Candidate);

        // Heartbeat from an older term should not cause stepdown
        node.ReceiveHeartbeat(term: 5);
        node.Role.ShouldBe(RaftRole.Candidate);
        node.Term.ShouldBe(11);
    }

    // Go: TestNRGSimple server/raft_test.go:35 — seven-node quorum
    [Theory]
    [InlineData(1, 1)]  // Single node: quorum = 1
    [InlineData(3, 2)]  // 3-node: quorum = 2
    [InlineData(5, 3)]  // 5-node: quorum = 3
    [InlineData(7, 4)]  // 7-node: quorum = 4
    public void Quorum_size_for_various_cluster_sizes(int clusterSize, int expectedQuorum)
    {
        var node = new RaftNode("n1");
        node.StartElection(clusterSize);

        // Self-vote = 1, need (expectedQuorum - 1) more
        for (int i = 0; i < expectedQuorum - 1; i++)
            node.ReceiveVote(new VoteResponse { Granted = true }, clusterSize);

        node.IsLeader.ShouldBeTrue();
    }
}