import threading
import time
import random
import logging
from enum import Enum
from dataclasses import dataclass, field
from typing import Optional

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(name)s] %(message)s')


class NodeState(Enum):
    FOLLOWER = "follower"
    CANDIDATE = "candidate"
    LEADER = "leader"


@dataclass
class LogEntry:
    term: int         # term when this entry was received by leader
    index: int        # 1-based position in the log
    command: str      # the state machine command (e.g., "SET x=5")


class RaftNode:
    """
    Simplified Raft implementation for educational purposes.
    In production, RPCs would be network calls; here they're in-process method calls.
    """

    def __init__(self, node_id: int, peers: list['RaftNode']):
        self.node_id = node_id
        self.peers = peers           # references to other nodes
        self.logger = logging.getLogger(f"Node-{node_id}")

        # --- Persistent state (must survive restarts in real systems) ---
        self.current_term = 0        # latest term seen
        self.voted_for: Optional[int] = None  # candidate voted for in current term
        self.log: list[LogEntry] = []         # log entries; index 0 = first entry

        # --- Volatile state (all nodes) ---
        self.commit_index = 0        # highest log entry known to be committed
        self.last_applied = 0        # highest log entry applied to state machine

        # --- Volatile state (leaders only, reinitialized after election) ---
        self.next_index: dict[int, int] = {}   # for each peer, next log index to send
        self.match_index: dict[int, int] = {}  # for each peer, highest replicated index

        # --- Node behavior ---
        self.state = NodeState.FOLLOWER
        self.leader_id: Optional[int] = None
        self.lock = threading.Lock()
        self.election_timeout = self._new_election_timeout()
        self.last_heartbeat = time.time()

        # State machine (just a dict for demo)
        self.state_machine: dict[str, str] = {}

    def _new_election_timeout(self) -> float:
        """
        Randomized timeout between 150-300ms prevents split votes.
        All followers use different timeouts so one becomes candidate first.
        """
        return random.uniform(0.15, 0.30)

    # =========================================================
    # PART A: Leader Election
    # =========================================================

    def start_election(self):
        """
        Called when a follower's election timer fires.
        Transition: Follower → Candidate
        """
        with self.lock:
            self.state = NodeState.CANDIDATE
            self.current_term += 1          # increment term before starting election
            self.voted_for = self.node_id   # vote for self
            term = self.current_term
            last_log_index = len(self.log)
            last_log_term = self.log[-1].term if self.log else 0
            self.logger.info(f"Starting election for term {term}")

        votes_received = 1  # self-vote
        votes_needed = (len(self.peers) + 1) // 2 + 1  # majority of total nodes

        for peer in self.peers:
            granted = peer.request_vote(
                term=term,
                candidate_id=self.node_id,
                last_log_index=last_log_index,
                last_log_term=last_log_term,
            )
            if granted:
                votes_received += 1
                if votes_received >= votes_needed:
                    self._become_leader()
                    return

        # Didn't win; revert to follower and wait for next timeout
        with self.lock:
            if self.state == NodeState.CANDIDATE:
                self.state = NodeState.FOLLOWER

    def request_vote(self, term: int, candidate_id: int,
                     last_log_index: int, last_log_term: int) -> bool:
        """
        RPC handler: another node is asking for our vote.
        Grant vote only if:
          1. candidate's term >= our current term
          2. we haven't voted for someone else this term
          3. candidate's log is at least as up-to-date as ours (Election Restriction)
        """
        with self.lock:
            # Step down if we see a higher term
            if term > self.current_term:
                self.current_term = term
                self.voted_for = None
                self.state = NodeState.FOLLOWER

            if term < self.current_term:
                return False  # stale term; reject

            # Check if we already voted this term
            if self.voted_for is not None and self.voted_for != candidate_id:
                return False

            # Election Restriction: candidate must be at least as up-to-date
            # "Up-to-date" = higher last term, OR same last term + longer log
            my_last_log_term = self.log[-1].term if self.log else 0
            my_last_log_index = len(self.log)

            candidate_is_up_to_date = (
                last_log_term > my_last_log_term or
                (last_log_term == my_last_log_term and last_log_index >= my_last_log_index)
            )

            if candidate_is_up_to_date:
                self.voted_for = candidate_id
                self.last_heartbeat = time.time()  # reset our timer
                self.logger.info(f"Voted for node {candidate_id} in term {term}")
                return True

            return False

    def _become_leader(self):
        """Transition to leader; initialize replication tracking."""
        with self.lock:
            self.state = NodeState.LEADER
            self.leader_id = self.node_id
            # nextIndex = last log index + 1 for all peers (optimistic)
            next_idx = len(self.log) + 1
            for peer in self.peers:
                self.next_index[peer.node_id] = next_idx
                self.match_index[peer.node_id] = 0
            self.logger.info(f"Became LEADER for term {self.current_term}")

    # =========================================================
    # PART B: Log Replication
    # =========================================================

    def client_request(self, command: str) -> bool:
        """
        Entry point for client commands.
        Only the leader should accept these; followers redirect.
        """
        with self.lock:
            if self.state != NodeState.LEADER:
                self.logger.warning(f"Not leader; redirecting to node {self.leader_id}")
                return False

            # Append to local log first
            entry = LogEntry(
                term=self.current_term,
                index=len(self.log) + 1,
                command=command,
            )
            self.log.append(entry)
            self.logger.info(f"Appended to log: {entry}")

        # Replicate to all followers; wait for quorum
        return self._replicate_log(entry.index)

    def _replicate_log(self, target_index: int) -> bool:
        """
        Send AppendEntries to all peers and wait for quorum acknowledgment.
        Returns True when the entry is safely committed.
        """
        acks = 1  # leader counts itself
        needed = (len(self.peers) + 1) // 2 + 1

        for peer in self.peers:
            success = self._send_append_entries(peer)
            if success:
                acks += 1
                if acks >= needed:
                    # Quorum reached — commit!
                    self._advance_commit_index(target_index)
                    return True

        return acks >= needed  # might still have quorum

    def _send_append_entries(self, peer: 'RaftNode') -> bool:
        """
        Send AppendEntries RPC to a single peer.
        Includes a consistency check: prevLogIndex/prevLogTerm must match.
        On failure, we decrement nextIndex and retry (back-tracking).
        """
        with self.lock:
            next_idx = self.next_index.get(peer.node_id, 1)
            prev_log_index = next_idx - 1
            prev_log_term = (
                self.log[prev_log_index - 1].term if prev_log_index > 0 else 0
            )
            entries = self.log[next_idx - 1:]  # entries to send
            term = self.current_term
            leader_commit = self.commit_index

        success, peer_term = peer.append_entries(
            term=term,
            leader_id=self.node_id,
            prev_log_index=prev_log_index,
            prev_log_term=prev_log_term,
            entries=entries,
            leader_commit=leader_commit,
        )

        with self.lock:
            if peer_term > self.current_term:
                # Discovered higher term; step down
                self.current_term = peer_term
                self.state = NodeState.FOLLOWER
                self.voted_for = None
                return False

            if success:
                # Update our tracking of this peer's log
                new_match = prev_log_index + len(entries)
                self.match_index[peer.node_id] = new_match
                self.next_index[peer.node_id] = new_match + 1
            else:
                # Consistency check failed; back off by 1 and retry next time
                self.next_index[peer.node_id] = max(1, self.next_index[peer.node_id] - 1)

        return success

    def append_entries(self, term: int, leader_id: int,
                       prev_log_index: int, prev_log_term: int,
                       entries: list[LogEntry], leader_commit: int):
        """
        RPC handler: leader is sending log entries (or heartbeat if entries=[]).
        Returns (success, my_current_term).
        """
        with self.lock:
            if term > self.current_term:
                self.current_term = term
                self.voted_for = None
                self.state = NodeState.FOLLOWER

            if term < self.current_term:
                return False, self.current_term  # stale leader; reject

            # Reset heartbeat timer — we heard from a valid leader
            self.last_heartbeat = time.time()
            self.leader_id = leader_id
            self.state = NodeState.FOLLOWER

            # Consistency check: does our log match up to prevLogIndex?
            if prev_log_index > 0:
                if len(self.log) < prev_log_index:
                    return False, self.current_term  # we're missing entries
                if self.log[prev_log_index - 1].term != prev_log_term:
                    # Conflict: delete from this index onward and retry
                    self.log = self.log[:prev_log_index - 1]
                    return False, self.current_term

            # Append new entries, overwriting any conflicts
            for i, entry in enumerate(entries):
                log_index = prev_log_index + i + 1
                if log_index <= len(self.log):
                    if self.log[log_index - 1].term != entry.term:
                        # Conflicting entry; truncate and append
                        self.log = self.log[:log_index - 1]
                        self.log.append(entry)
                else:
                    self.log.append(entry)

            # Advance commit index if leader says we can
            if leader_commit > self.commit_index:
                self.commit_index = min(leader_commit, len(self.log))
                self._apply_committed_entries()

            return True, self.current_term

    def _advance_commit_index(self, new_index: int):
        """
        Leader advances commit_index once a quorum has replicated the entry.
        Critical safety rule: only commit entries from the CURRENT term.
        (Prevents the 'Figure 8' bug from the Raft paper.)
        """
        with self.lock:
            if (new_index > self.commit_index and
                    self.log[new_index - 1].term == self.current_term):
                self.commit_index = new_index
                self._apply_committed_entries()

    def _apply_committed_entries(self):
        """Apply all committed but not-yet-applied entries to the state machine."""
        while self.last_applied < self.commit_index:
            self.last_applied += 1
            entry = self.log[self.last_applied - 1]
            # Parse "SET key=value" commands
            if entry.command.startswith("SET "):
                parts = entry.command[4:].split("=", 1)
                if len(parts) == 2:
                    self.state_machine[parts[0].strip()] = parts[1].strip()
            self.logger.info(
                f"Applied entry {self.last_applied}: {entry.command} "
                f"→ state={self.state_machine}"
            )


# =========================================================
# DEMO: 3-node Raft cluster
# =========================================================

def demo_raft_cluster():
    """
    Create a 3-node cluster, elect a leader, and replicate some commands.
    This is in-process (no real network), so we manually trigger election.
    """
    # Create nodes (no peers yet)
    nodes = [RaftNode(i, []) for i in range(3)]

    # Wire up peers (each node knows about the other two)
    for i, node in enumerate(nodes):
        node.peers = [n for n in nodes if n.node_id != i]

    # Manually start election on node 0
    print("\n=== Starting election on Node 0 ===")
    nodes[0].start_election()

    # Find who became leader
    leader = next((n for n in nodes if n.state == NodeState.LEADER), None)
    if leader:
        print(f"\n=== Leader elected: Node {leader.node_id} ===")

        # Submit commands
        leader.client_request("SET name=Alice")
        leader.client_request("SET age=30")

        # Check state machine on all nodes
        print("\n=== State machines after replication ===")
        for node in nodes:
            print(f"  Node {node.node_id}: {node.state_machine}")
    else:
        print("No leader elected (split vote?)")


if __name__ == "__main__":
    demo_raft_cluster()