"""Stdlib-only tests (no pytest): plain asserts in functions called from a
__main__ block. Run: `python3 test_raft.py`.

The headline test pins the whole Figure 8 scenario -- b@2 on a majority of 3
at step (c), the naive rule calling it committed and the real rule refusing,
and the entry gone from all five nodes by step (e) (commentary.html section 6).
The rest are unit checks on the election restriction, the log-truncation path
and the commit rule itself.
"""

import contextlib
import io

import demo
from raft import Cluster, Entry, Node, NAIVE, RAFT


def quiet(fn, *args, **kwargs):
    """Run something noisy and return (result, captured stdout)."""
    buf = io.StringIO()
    with contextlib.redirect_stdout(buf):
        result = fn(*args, **kwargs)
    return result, buf.getvalue()


# ---- the headline ------------------------------------------------------

def test_naive_commits_an_entry_that_a_later_leader_erases():
    """The aha. Identical schedule, one clause different, opposite verdicts."""
    (naive, naive_claimed), _ = quiet(demo.figure8, NAIVE, verbose=False)
    (raft, raft_claimed), _ = quiet(demo.figure8, RAFT, verbose=False)

    assert naive_claimed is True, "naive rule should have claimed b@2 committed"
    assert raft_claimed is False, "raft rule must refuse to commit b@2"

    # ...and the entry the naive rule vouched for exists nowhere afterwards.
    assert naive.holders(2, "b") == [], naive.holders(2, "b")
    assert raft.holders(2, "b") == [], raft.holders(2, "b")
    for node in naive.nodes.values():
        assert [e.cmd for e in node.log] == ["a", "c"], node.log


def test_the_majority_at_step_c_is_real():
    """b@2 is on exactly 3 of 5 -- not a near-majority the rule is dodging."""
    for rule in (NAIVE, RAFT):
        c = Cluster(5, rule)
        c.elect(1)
        c.nodes[1].client("a")
        c.deliver()
        c.deliver()
        c.elect(1)
        c.nodes[1].client("b")
        c.deliver(only={1, 2})
        c.crash(1)
        c.elect(5)
        c.nodes[5].client("c")
        c.deliver(only={5})
        c.crash(5)
        c.restart(1)
        assert c.elect(1, only={1, 2, 3}) == 4
        c.deliver(only={1, 2, 3})
        assert c.holders(2, "b") == [1, 2, 3], c.holders(2, "b")
        leader = c.nodes[1]
        stored = 1 + sum(1 for p in leader.peers if leader.match_index[p] >= 2)
        assert stored == 3, stored
        assert stored * 2 > 5, "3 of 5 is a majority"
        assert leader.log[1].term == 2 and leader.term == 4
        # The naive rule counts 3 replicas and commits. The real rule commits
        # nothing at all: both entries in the log predate term 4.
        assert leader.commit_index == (2 if rule == NAIVE else 0)


def test_the_two_rules_differ_only_in_the_term_clause():
    """Same match_index, same log, same term: only `rule` decides."""
    def leader_with(rule):
        node = Node(1, [2, 3, 4, 5], rule)
        node.out = []
        node.role = "leader"
        node.term = 4
        node.log = [Entry(1, "a"), Entry(2, "b")]
        node.match_index = {2: 2, 3: 2, 4: 0, 5: 0}
        node.advance_commit()
        return node.commit_index

    assert leader_with(NAIVE) == 2
    # RAFT commits NOTHING -- not index 2, and not index 1 either, because that
    # entry is from term 1. A fresh leader is stuck until it writes something.
    assert leader_with(RAFT) == 0

    # ...which is exactly what a no-op entry of the leader's own term fixes.
    node = Node(1, [2, 3, 4, 5], RAFT)
    node.out = []
    node.role = "leader"
    node.term = 4
    node.log = [Entry(1, "a"), Entry(2, "b"), Entry(4, "no-op")]
    node.match_index = {2: 3, 3: 3, 4: 0, 5: 0}
    node.advance_commit()
    assert node.commit_index == 3, "the no-op commits, and 1 and 2 ride along"


# ---- the boundary ------------------------------------------------------

def test_a_current_term_entry_carries_the_old_one_with_it():
    """Section 6's boundary: one term-4 entry and b@2 becomes unkillable."""
    c = Cluster(5, RAFT)
    c.elect(1)
    c.nodes[1].client("a")
    c.deliver()
    c.deliver()
    c.elect(1)
    c.nodes[1].client("b")
    c.deliver(only={1, 2})
    c.crash(1)
    c.elect(5)
    c.nodes[5].client("c")
    c.deliver(only={5})
    c.crash(5)
    c.restart(1)
    c.elect(1, only={1, 2, 3})
    c.nodes[1].client("d")                    # the one difference
    c.deliver(only={1, 2, 3})
    c.deliver(only={1, 2, 3})

    leader = c.nodes[1]
    assert leader.log[2].term == leader.term == 4
    assert leader.commit_index == 3, leader.commit_index   # index 2 rides along

    c.crash(1)
    c.restart(5)
    c.elect(5, only={2, 3, 4, 5})
    assert c.nodes[5].role != "leader", "S5 must never win once index 3 commits"
    assert c.holders(2, "b") == [1, 2, 3], c.holders(2, "b")


# ---- the negative result -----------------------------------------------

def test_length_only_vote_comparison_does_not_save_figure_8():
    """Section 7's negative result. Weakening the election restriction to
    compare log LENGTH instead of (term, index) changes nothing here: the
    commit rule is what saves Figure 8, not the vote check."""
    original = Node.on_vote

    def by_length(self, m):
        idx, _ = self.last()
        up_to_date = m["last_index"] >= idx        # term ignored
        granted = (m["term"] == self.term
                   and self.voted_for in (None, m["src"])
                   and up_to_date)
        if granted:
            self.voted_for = m["src"]
        self.send(m["src"], type="vote_reply", term=self.term, granted=granted)

    Node.on_vote = by_length
    try:
        (c, claimed), _ = quiet(demo.figure8, RAFT, verbose=False)
    finally:
        Node.on_vote = original
    assert claimed is False
    assert c.holders(2, "b") == [], "S5 still wins term 5 and still truncates"
    assert [e.cmd for e in c.nodes[5].log] == ["a", "c"]


# ---- unit checks -------------------------------------------------------

def test_election_restriction_compares_term_before_index():
    """A longer log with an older final term loses to a shorter, newer one."""
    voter = Node(2, [1], RAFT)
    voter.out = []
    voter.term = 5
    voter.log = [Entry(1, "a"), Entry(2, "b"), Entry(2, "x")]   # 3 entries
    voter.recv(dict(src=1, dst=2, type="vote", term=5, last_index=2, last_term=3))
    assert voter.out[-1]["granted"] is True, "term 3 beats term 2 despite being shorter"

    voter2 = Node(2, [1], RAFT)
    voter2.out = []
    voter2.term = 5
    voter2.log = [Entry(1, "a"), Entry(3, "c")]
    voter2.recv(dict(src=1, dst=2, type="vote", term=5, last_index=9, last_term=2))
    assert voter2.out[-1]["granted"] is False, "a longer log with an older term loses"


def test_a_node_votes_at_most_once_per_term():
    voter = Node(2, [1, 3], RAFT)
    voter.out = []
    voter.recv(dict(src=1, dst=2, type="vote", term=1, last_index=0, last_term=0))
    assert voter.out[-1]["granted"] is True and voter.voted_for == 1
    voter.recv(dict(src=3, dst=2, type="vote", term=1, last_index=0, last_term=0))
    assert voter.out[-1]["granted"] is False, "second candidate in term 1 is denied"
    # A higher term resets the vote, which is why terms must survive a crash.
    voter.recv(dict(src=3, dst=2, type="vote", term=2, last_index=0, last_term=0))
    assert voter.out[-1]["granted"] is True and voter.voted_for == 3


def test_stale_candidate_loses_then_wins_on_the_retry():
    """`elect` burning terms is real behaviour: S1 at term 2 cannot beat the
    term-3 nodes on its first attempt, learns the term, and wins on the second."""
    c = Cluster(5, RAFT)
    c.elect(1)
    c.nodes[1].client("a")
    c.deliver()
    c.crash(1)
    c.elect(5)
    assert c.nodes[5].term == 2
    c.crash(5)
    c.restart(1)
    assert c.nodes[1].term == 1
    c.nodes[1].campaign()                      # term 2 -- ties, cannot win
    c.deliver(only={1, 2, 3})
    assert c.nodes[1].role != "leader"
    assert c.nodes[1].term == 2
    c.nodes[1].campaign()                      # term 3 -- wins
    c.deliver(only={1, 2, 3})
    assert c.nodes[1].role == "leader" and c.nodes[1].term == 3


def test_next_index_walks_back_until_the_logs_agree():
    """A new leader assumes followers match, and the rejections correct it."""
    c = Cluster(5, RAFT)
    c.nodes[1].log = [Entry(1, "a"), Entry(1, "b"), Entry(1, "c")]
    c.nodes[2].log = [Entry(1, "a")]
    for i in (1, 2):
        c.nodes[i].term = 1
    c.nodes[1].role = "leader"
    c.nodes[1].next_index = {p: 4 for p in c.nodes[1].peers}
    c.nodes[1].match_index = {p: 0 for p in c.nodes[1].peers}
    c.nodes[1].send_append(2)
    c.deliver(only={1, 2})
    assert [e.cmd for e in c.nodes[2].log] == ["a", "b", "c"]
    assert c.nodes[1].match_index[2] == 3


def test_a_follower_truncates_a_conflicting_suffix():
    """The line that erases b@2: the leader's log wins, unconditionally."""
    c = Cluster(5, RAFT)
    c.nodes[1].log = [Entry(1, "a"), Entry(3, "c")]
    c.nodes[2].log = [Entry(1, "a"), Entry(2, "b")]
    for i in (1, 2):
        c.nodes[i].term = 3
    c.nodes[1].role = "leader"
    c.nodes[1].next_index = {p: 3 for p in c.nodes[1].peers}
    c.nodes[1].match_index = {p: 0 for p in c.nodes[1].peers}
    c.nodes[1].send_append(2)
    c.deliver(only={1, 2})
    assert [repr(e) for e in c.nodes[2].log] == ["a@1", "c@3"]


def test_a_crash_keeps_the_log_term_and_vote():
    c = Cluster(5, RAFT)
    c.elect(1)
    c.nodes[1].client("a")
    c.deliver()
    before = (c.nodes[1].term, c.nodes[1].voted_for, [repr(e) for e in c.nodes[1].log])
    c.crash(1)
    c.restart(1)
    after = (c.nodes[1].term, c.nodes[1].voted_for, [repr(e) for e in c.nodes[1].log])
    assert before == after
    assert c.nodes[1].role == "follower", "volatile leadership is lost"


def test_messages_outside_the_partition_are_lost_not_queued():
    c = Cluster(5, RAFT)
    c.elect(1)
    c.nodes[1].client("a")
    c.deliver(only={1, 2})
    assert c.net == [], "undelivered messages are dropped, not held"
    assert len(c.nodes[2].log) == 1 and c.nodes[3].log == []


def test_demo_output_is_byte_identical_across_runs():
    """No clock, no RNG, no threads: the schedule is the whole input."""
    _, first = quiet(demo.run_figure8)
    _, second = quiet(demo.run_figure8)
    assert first == second
    _, b1 = quiet(demo.boundary)
    _, b2 = quiet(demo.boundary)
    assert b1 == b2


TESTS = [
    test_naive_commits_an_entry_that_a_later_leader_erases,
    test_the_majority_at_step_c_is_real,
    test_the_two_rules_differ_only_in_the_term_clause,
    test_a_current_term_entry_carries_the_old_one_with_it,
    test_length_only_vote_comparison_does_not_save_figure_8,
    test_election_restriction_compares_term_before_index,
    test_a_node_votes_at_most_once_per_term,
    test_stale_candidate_loses_then_wins_on_the_retry,
    test_next_index_walks_back_until_the_logs_agree,
    test_a_follower_truncates_a_conflicting_suffix,
    test_a_crash_keeps_the_log_term_and_vote,
    test_messages_outside_the_partition_are_lost_not_queued,
    test_demo_output_is_byte_identical_across_runs,
]


if __name__ == "__main__":
    failed = 0
    for test in TESTS:
        try:
            test()
        except AssertionError as exc:
            failed += 1
            print(f"FAIL  {test.__name__}: {exc}")
        else:
            print(f"PASS  {test.__name__}")
    if failed:
        print(f"\n{failed} of {len(TESTS)} tests FAILED")
        raise SystemExit(1)
    print(f"\nAll {len(TESTS)} tests PASSED")