"""Minimal Raft consensus for tracker shard assignments.

Only shard-assignment commands (register/deregister) go through the log.
Node liveness (heartbeats) is handled separately via CRDT gossip — these
are high-frequency writes where eventual consistency is fine.

Election timeout: random 150–300 ms (tight, suits in-process tests).
Leader heartbeat interval: 50 ms.
"""

from __future__ import annotations

import json
import random
import threading
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Callable


@dataclass
class LogEntry:
    term: int
    command: str   # "register" | "deregister"
    payload: dict


class RaftNode:
    """Single Raft participant.

    ``apply_fn(command, payload)`` is called (under no external lock) when an
    entry is committed.  Implementors must apply the command atomically.
    """

    ELECTION_MIN = 0.15   # seconds
    ELECTION_MAX = 0.30
    HB_INTERVAL  = 0.05   # leader heartbeat interval

    # Role constants
    FOLLOWER  = "follower"
    CANDIDATE = "candidate"
    LEADER    = "leader"

    def __init__(
        self,
        self_url: str,
        peers: list[str],
        apply_fn: Callable[[str, dict], None],
    ) -> None:
        self.self_url = self_url
        self.peers = list(peers)
        self._apply_fn = apply_fn

        self._lock = threading.Lock()
        self.role: str = self.FOLLOWER
        self.current_term: int = 0
        self.voted_for: str | None = None
        self.log: list[LogEntry] = []
        self.commit_index: int = -1
        self.last_applied: int = -1
        self.leader_url: str | None = None

        # Leader bookkeeping per peer
        self._next_index:  dict[str, int] = {}
        self._match_index: dict[str, int] = {}

        self._election_deadline: float = self._fresh_deadline()
        self._running = False

    # ------------------------------------------------------------------ start/stop

    def start(self) -> None:
        self._running = True
        threading.Thread(target=self._tick_loop, daemon=True, name="raft-tick").start()

    def stop(self) -> None:
        self._running = False

    # ------------------------------------------------------------------ public API

    @property
    def is_leader(self) -> bool:
        with self._lock:
            return self.role == self.LEADER

    def leader(self) -> str | None:
        with self._lock:
            return self.leader_url

    def status(self) -> dict:
        with self._lock:
            return {
                "role": self.role,
                "term": self.current_term,
                "leader": self.leader_url,
                "log_length": len(self.log),
                "commit_index": self.commit_index,
            }

    def propose(self, command: str, payload: dict) -> bool:
        """Leader: append and replicate an entry. Returns True when committed.

        Blocks until majority replication or failure.  Must be called only on
        the leader; returns False immediately if this node is not the leader.
        """
        with self._lock:
            if self.role != self.LEADER:
                return False
            entry = LogEntry(self.current_term, command, payload)
            self.log.append(entry)
            entry_index = len(self.log) - 1
            term = self.current_term

        self._replicate_to_peers(term)

        with self._lock:
            return self.commit_index >= entry_index

    # ------------------------------------------------------------------ RPC handlers

    def handle_request_vote(self, req: dict) -> dict:
        with self._lock:
            term        = int(req["term"])
            candidate   = req["candidate_url"]
            last_li     = int(req["last_log_index"])
            last_lt     = int(req["last_log_term"])

            if term > self.current_term:
                self._step_down(term)

            if term < self.current_term:
                return {"term": self.current_term, "vote_granted": False}

            my_li = len(self.log) - 1
            my_lt = self.log[-1].term if self.log else 0
            log_ok = (last_lt > my_lt) or (last_lt == my_lt and last_li >= my_li)

            if (self.voted_for is None or self.voted_for == candidate) and log_ok:
                self.voted_for = candidate
                self._reset_deadline()
                return {"term": self.current_term, "vote_granted": True}
            return {"term": self.current_term, "vote_granted": False}

    def handle_append_entries(self, req: dict) -> dict:
        with self._lock:
            term         = int(req["term"])
            leader_url   = req["leader_url"]
            prev_li      = int(req["prev_log_index"])
            prev_lt      = int(req["prev_log_term"])
            entries_raw  = req.get("entries", [])
            ldr_commit   = int(req["leader_commit"])

            if term > self.current_term:
                self._step_down(term)

            if term < self.current_term:
                return {"term": self.current_term, "success": False}

            # Valid AppendEntries from current leader
            self._reset_deadline()
            self.leader_url = leader_url
            if self.role == self.CANDIDATE:
                self.role = self.FOLLOWER

            # Consistency check on prev entry
            if prev_li >= 0:
                if prev_li >= len(self.log):
                    return {"term": self.current_term, "success": False}
                if self.log[prev_li].term != prev_lt:
                    self.log = self.log[:prev_li]
                    return {"term": self.current_term, "success": False}

            # Append new entries (detect and overwrite conflicts)
            for i, raw in enumerate(entries_raw):
                idx = prev_li + 1 + i
                entry = LogEntry(int(raw["term"]), raw["command"], raw["payload"])
                if idx < len(self.log):
                    if self.log[idx].term != entry.term:
                        self.log = self.log[:idx]
                        self.log.append(entry)
                else:
                    self.log.append(entry)

            if ldr_commit > self.commit_index:
                self.commit_index = min(ldr_commit, len(self.log) - 1)
                self._apply_up_to_commit()

            return {"term": self.current_term, "success": True}

    # ------------------------------------------------------------------ internals

    def _tick_loop(self) -> None:
        while self._running:
            time.sleep(0.01)
            with self._lock:
                role     = self.role
                deadline = self._election_deadline

            if role == self.LEADER:
                self._send_heartbeats()
                time.sleep(self.HB_INTERVAL)
            elif time.monotonic() > deadline:
                self._start_election()

    def _send_heartbeats(self) -> None:
        with self._lock:
            term         = self.current_term
            ldr_commit   = self.commit_index
            log_snapshot = list(self.log)

        for peer in self.peers:
            try:
                self._send_append_entries(peer, term, ldr_commit, log_snapshot)
            except Exception:
                pass

    def _replicate_to_peers(self, term: int) -> None:
        """Send AppendEntries to all peers and update commit_index on majority ack."""
        with self._lock:
            ldr_commit   = self.commit_index
            log_snapshot = list(self.log)

        acks = 1  # leader counts as 1
        for peer in self.peers:
            try:
                ok = self._send_append_entries(peer, term, ldr_commit, log_snapshot)
                if ok:
                    acks += 1
            except Exception:
                pass

        # Advance commit if majority replicated
        with self._lock:
            if self.role != self.LEADER or self.current_term != term:
                return
            majority_index = len(self.log) - 1
            while majority_index > self.commit_index:
                if self.log[majority_index].term == self.current_term:
                    count = 1 + sum(
                        1 for p in self.peers
                        if self._match_index.get(p, -1) >= majority_index
                    )
                    if count > (len(self.peers) + 1) / 2:
                        self.commit_index = majority_index
                        self._apply_up_to_commit()
                        break
                majority_index -= 1

    def _send_append_entries(
        self, peer: str, term: int, ldr_commit: int, log_snapshot: list[LogEntry]
    ) -> bool:
        with self._lock:
            next_idx = self._next_index.get(peer, len(log_snapshot))

        prev_li = next_idx - 1
        prev_lt = log_snapshot[prev_li].term if 0 <= prev_li < len(log_snapshot) else 0
        entries = [
            {"term": e.term, "command": e.command, "payload": e.payload}
            for e in log_snapshot[next_idx:]
        ]

        body = json.dumps({
            "term": term,
            "leader_url": self.self_url,
            "prev_log_index": prev_li,
            "prev_log_term": prev_lt,
            "entries": entries,
            "leader_commit": ldr_commit,
        }).encode()
        req = urllib.request.Request(
            f"{peer}/v1/raft/append",
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=0.15) as r:
            resp = json.loads(r.read())

        with self._lock:
            if resp.get("success"):
                new_match = len(log_snapshot) - 1
                self._match_index[peer] = max(self._match_index.get(peer, -1), new_match)
                self._next_index[peer]  = new_match + 1
                return True
            else:
                if resp.get("term", 0) > self.current_term:
                    self._step_down(resp["term"])
                else:
                    self._next_index[peer] = max(0, self._next_index.get(peer, 1) - 1)
                return False

    def _start_election(self) -> None:
        with self._lock:
            self.current_term += 1
            self.role = self.CANDIDATE
            self.voted_for = self.self_url
            term  = self.current_term
            my_li = len(self.log) - 1
            my_lt = self.log[-1].term if self.log else 0
            self._reset_deadline()

        votes = 1
        for peer in self.peers:
            try:
                body = json.dumps({
                    "term": term,
                    "candidate_url": self.self_url,
                    "last_log_index": my_li,
                    "last_log_term": my_lt,
                }).encode()
                req = urllib.request.Request(
                    f"{peer}/v1/raft/vote",
                    data=body,
                    headers={"Content-Type": "application/json"},
                    method="POST",
                )
                with urllib.request.urlopen(req, timeout=0.1) as r:
                    resp = json.loads(r.read())
                if resp.get("vote_granted"):
                    votes += 1
                elif resp.get("term", 0) > term:
                    with self._lock:
                        self._step_down(resp["term"])
                    return
            except Exception:
                pass

        with self._lock:
            if self.role == self.CANDIDATE and self.current_term == term:
                total = len(self.peers) + 1
                if votes > total / 2:
                    self._become_leader()
                else:
                    self.role = self.FOLLOWER
                    self._reset_deadline()

    def _become_leader(self) -> None:
        """Must be called with _lock held."""
        self.role = self.LEADER
        self.leader_url = self.self_url
        for peer in self.peers:
            self._next_index[peer]  = len(self.log)
            self._match_index[peer] = -1
        print(f"[raft] {self.self_url} became leader (term {self.current_term})", flush=True)

    def _step_down(self, new_term: int) -> None:
        """Must be called with _lock held."""
        self.current_term = new_term
        self.role = self.FOLLOWER
        self.voted_for = None
        self._reset_deadline()

    def _apply_up_to_commit(self) -> None:
        """Must be called with _lock held."""
        while self.last_applied < self.commit_index:
            self.last_applied += 1
            entry = self.log[self.last_applied]
            try:
                self._apply_fn(entry.command, entry.payload)
            except Exception:
                pass

    def _reset_deadline(self) -> None:
        self._election_deadline = self._fresh_deadline()

    @staticmethod
    def _fresh_deadline() -> float:
        return time.monotonic() + random.uniform(
            RaftNode.ELECTION_MIN, RaftNode.ELECTION_MAX
        )
