Leader Election and Coordination.

# Fencing Token Implementation

from dataclasses import dataclass
from typing import Optional, Any
from datetime import datetime, timedelta, timezone
import threading
import time


@dataclass
class LeaderLease:
    """A leader's lease with fencing token."""
    leader_id: str
    token: int
    expires_at: datetime
    
    def is_valid(self) -> bool:
        """Check if lease is still valid."""
        return datetime.now(timezone.utc) < self.expires_at


class FencedResource:
    """
    A resource protected by fencing tokens.
    
    Rejects operations from stale leaders.
    """
    
    def __init__(self, resource_id: str):
        self.resource_id = resource_id
        self.highest_token_seen = 0
        self.lock = threading.Lock()
        self.data: dict = {}
    
    def write(self, key: str, value: Any, token: int) -> bool:
        """
        Write to resource with fencing token.
        
        Returns True if write was accepted.
        """
        with self.lock:
            if token < self.highest_token_seen:
                # Stale leader - reject
                return False
            
            self.highest_token_seen = token
            self.data[key] = value
            return True
    
    def read(self, key: str) -> Optional[Any]:
        """Read from resource."""
        return self.data.get(key)


class FencedLeaderClient:
    """
    Client that uses fencing tokens for leader operations.
    """
    
    def __init__(self, client_id: str):
        self.client_id = client_id
        self.current_lease: Optional[LeaderLease] = None
    
    def set_lease(self, lease: LeaderLease):
        """Set the current lease (from election)."""
        self.current_lease = lease
    
    def execute_as_leader(
        self,
        resource: FencedResource,
        key: str,
        value: Any
    ) -> bool:
        """
        Execute a write operation as leader.
        
        Includes fencing token validation.
        """
        if not self.current_lease:
            raise RuntimeError("Not a leader - no lease")
        
        if not self.current_lease.is_valid():
            raise RuntimeError("Lease expired")
        
        # Include fencing token with request
        success = resource.write(
            key,
            value,
            self.current_lease.token
        )
        
        if not success:
            # We're a zombie leader - stop operations
            self.current_lease = None
            raise RuntimeError("Fenced - no longer leader")
        
        return True


# Example: Database with fencing
class FencedDatabaseClient:
    """
    Database client that enforces fencing tokens.
    
    The check must live in the database, atomically with the write:
    a counter kept in the client is per-process, so a zombie leader's
    own client would happily accept its stale token.
    """
    
    def __init__(self, db_connection):
        self.db = db_connection
    
    async def write(self, resource_id: str, value: Any, token: int) -> Any:
        """Write only if no newer token has been seen for this resource."""
        result = await self.db.execute(
            """
            UPDATE resources SET value = $2, fence_token = $3
            WHERE id = $1 AND fence_token <= $3
            """,
            resource_id, value, token
        )
        if result == "UPDATE 0":
            raise FencingError(f"Stale token {token} for {resource_id}")
        return result


# Example: Distributed lock with fencing
class FencedLock:
    """
    Distributed lock that provides fencing tokens.
    """
    
    def __init__(self, lock_service, lock_name: str):
        self.service = lock_service
        self.lock_name = lock_name
        self.token: Optional[int] = None
    
    async def acquire(self, ttl_seconds: int = 30) -> int:
        """
        Acquire lock and return fencing token.
        
        The token should be passed to all protected resources.
        """
        result = await self.service.acquire_lock(
            self.lock_name,
            ttl_seconds=ttl_seconds
        )
        
        self.token = result.fencing_token
        return self.token
    
    async def release(self):
        """Release the lock."""
        if self.token:
            await self.service.release_lock(self.lock_name, self.token)
            self.token = None
    
    def get_token(self) -> Optional[int]:
        """Get current fencing token."""
        return self.token

# Simplified Raft Implementation

import asyncio
import random
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Set
from enum import Enum
from datetime import datetime, timedelta, timezone
import logging

logger = logging.getLogger(__name__)


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


@dataclass
class LogEntry:
    """Entry in the Raft log."""
    term: int
    index: int
    command: str
    data: dict


@dataclass
class RequestVoteRequest:
    term: int
    candidate_id: str
    last_log_index: int
    last_log_term: int


@dataclass
class RequestVoteResponse:
    term: int
    vote_granted: bool


@dataclass
class AppendEntriesRequest:
    term: int
    leader_id: str
    prev_log_index: int
    prev_log_term: int
    entries: List[LogEntry]
    leader_commit: int


@dataclass
class AppendEntriesResponse:
    term: int
    success: bool


class RaftNode:
    """
    Simplified Raft consensus node.
    
    This implementation focuses on leader election.
    Log replication is simplified for clarity.
    """
    
    def __init__(
        self,
        node_id: str,
        peers: List[str],
        rpc_client,  # Client to communicate with peers
        election_timeout_ms: tuple = (150, 300),  # Random range
        heartbeat_interval_ms: int = 50
    ):
        self.node_id = node_id
        self.peers = peers
        self.rpc = rpc_client
        self.election_timeout_range = election_timeout_ms
        self.heartbeat_interval = heartbeat_interval_ms / 1000
        
        # Persistent state
        self.current_term = 0
        self.voted_for: Optional[str] = None
        self.log: List[LogEntry] = []
        
        # Volatile state
        self.state = NodeState.FOLLOWER
        self.commit_index = 0
        self.last_applied = 0
        
        # Leader state
        self.next_index: Dict[str, int] = {}
        self.match_index: Dict[str, int] = {}
        
        # Timing
        self.last_heartbeat = datetime.now(timezone.utc)
        self.election_timeout = self._random_election_timeout()
        
        # Control
        self.running = False
    
    def _random_election_timeout(self) -> float:
        """Get random election timeout in seconds."""
        min_ms, max_ms = self.election_timeout_range
        return random.randint(min_ms, max_ms) / 1000
    
    async def start(self):
        """Start the Raft node."""
        self.running = True
        asyncio.create_task(self._run_loop())
    
    async def stop(self):
        """Stop the Raft node."""
        self.running = False
    
    async def _run_loop(self):
        """Main loop for Raft node."""
        while self.running:
            if self.state == NodeState.LEADER:
                await self._leader_loop()
            else:
                await self._follower_candidate_loop()
            
            await asyncio.sleep(0.01)  # Small delay
    
    async def _follower_candidate_loop(self):
        """Loop for follower/candidate state."""
        now = datetime.now(timezone.utc)
        time_since_heartbeat = (now - self.last_heartbeat).total_seconds()
        
        if time_since_heartbeat > self.election_timeout:
            # Election timeout - start election
            await self._start_election()
    
    async def _leader_loop(self):
        """Loop for leader state - send heartbeats."""
        await self._send_heartbeats()
        await asyncio.sleep(self.heartbeat_interval)
    
    async def _start_election(self):
        """Start a leader election."""
        self.state = NodeState.CANDIDATE
        self.current_term += 1
        self.voted_for = self.node_id
        self.election_timeout = self._random_election_timeout()
        self.last_heartbeat = datetime.now(timezone.utc)
        
        logger.info(f"Node {self.node_id} starting election for term {self.current_term}")
        
        # Request votes from all peers
        votes_received = 1  # Vote for self
        
        last_log_index = len(self.log) - 1 if self.log else -1
        last_log_term = self.log[last_log_index].term if self.log else 0
        
        request = RequestVoteRequest(
            term=self.current_term,
            candidate_id=self.node_id,
            last_log_index=last_log_index,
            last_log_term=last_log_term
        )
        
        # Send to all peers in parallel
        tasks = [
            self._request_vote(peer, request)
            for peer in self.peers
        ]
        
        responses = await asyncio.gather(*tasks, return_exceptions=True)
        
        for response in responses:
            if isinstance(response, Exception):
                continue
            
            if response.term > self.current_term:
                # Discovered higher term - step down
                self._step_down(response.term)
                return
            
            if response.vote_granted:
                votes_received += 1
        
        # Check if we won
        majority = (len(self.peers) + 1) // 2 + 1
        
        if votes_received >= majority and self.state == NodeState.CANDIDATE:
            self._become_leader()
    
    async def _request_vote(self, peer: str, request: RequestVoteRequest) -> RequestVoteResponse:
        """Request vote from a peer."""
        try:
            response = await self.rpc.request_vote(peer, request)
            return response
        except Exception as e:
            logger.warning(f"Failed to request vote from {peer}: {e}")
            raise
    
    def _become_leader(self):
        """Transition to leader state."""
        self.state = NodeState.LEADER
        logger.info(f"Node {self.node_id} became leader for term {self.current_term}")
        
        # Initialize leader state
        next_idx = len(self.log)
        for peer in self.peers:
            self.next_index[peer] = next_idx
            self.match_index[peer] = 0
    
    def _step_down(self, term: int):
        """Step down to follower state."""
        self.current_term = term
        self.state = NodeState.FOLLOWER
        self.voted_for = None
        logger.info(f"Node {self.node_id} stepping down, term={term}")
    
    async def _send_heartbeats(self):
        """Send heartbeats to all peers."""
        tasks = [
            self._send_append_entries(peer)
            for peer in self.peers
        ]
        await asyncio.gather(*tasks, return_exceptions=True)
    
    async def _send_append_entries(self, peer: str):
        """Send AppendEntries to a peer."""
        prev_log_index = self.next_index.get(peer, 0) - 1
        prev_log_term = 0
        if prev_log_index >= 0 and prev_log_index < len(self.log):
            prev_log_term = self.log[prev_log_index].term
        
        # Get entries to send
        entries = self.log[self.next_index.get(peer, 0):]
        
        request = AppendEntriesRequest(
            term=self.current_term,
            leader_id=self.node_id,
            prev_log_index=prev_log_index,
            prev_log_term=prev_log_term,
            entries=entries,
            leader_commit=self.commit_index
        )
        
        try:
            response = await self.rpc.append_entries(peer, request)
            
            if response.term > self.current_term:
                self._step_down(response.term)
                return
            
            if response.success:
                # Update next_index and match_index
                self.next_index[peer] = prev_log_index + len(entries) + 1
                self.match_index[peer] = self.next_index[peer] - 1
            else:
                # Decrement next_index and retry
                self.next_index[peer] = max(0, self.next_index.get(peer, 1) - 1)
                
        except Exception as e:
            logger.warning(f"Failed to send AppendEntries to {peer}: {e}")
    
    # ==========================================================================
    # RPC Handlers
    # ==========================================================================
    
    async def handle_request_vote(self, request: RequestVoteRequest) -> RequestVoteResponse:
        """Handle RequestVote RPC from a candidate."""
        
        # If request term is higher, update our term
        if request.term > self.current_term:
            self._step_down(request.term)
        
        # Deny vote if our term is higher
        if request.term < self.current_term:
            return RequestVoteResponse(
                term=self.current_term,
                vote_granted=False
            )
        
        # Check if we can vote for this candidate
        can_vote = (
            self.voted_for is None or
            self.voted_for == request.candidate_id
        )
        
        # Check if candidate's log is at least as up-to-date
        last_log_index = len(self.log) - 1 if self.log else -1
        last_log_term = self.log[last_log_index].term if self.log else 0
        
        log_ok = (
            request.last_log_term > last_log_term or
            (request.last_log_term == last_log_term and
             request.last_log_index >= last_log_index)
        )
        
        if can_vote and log_ok:
            self.voted_for = request.candidate_id
            self.last_heartbeat = datetime.now(timezone.utc)  # Reset election timer
            
            return RequestVoteResponse(
                term=self.current_term,
                vote_granted=True
            )
        
        return RequestVoteResponse(
            term=self.current_term,
            vote_granted=False
        )
    
    async def handle_append_entries(self, request: AppendEntriesRequest) -> AppendEntriesResponse:
        """Handle AppendEntries RPC from leader."""
        
        # If request term is higher, update our term
        if request.term > self.current_term:
            self._step_down(request.term)
        
        # Reject if term is lower
        if request.term < self.current_term:
            return AppendEntriesResponse(
                term=self.current_term,
                success=False
            )
        
        # Valid AppendEntries from leader - reset election timer
        self.last_heartbeat = datetime.now(timezone.utc)
        self.state = NodeState.FOLLOWER
        
        # Check log consistency
        if request.prev_log_index >= 0:
            if request.prev_log_index >= len(self.log):
                return AppendEntriesResponse(
                    term=self.current_term,
                    success=False
                )
            
            if self.log[request.prev_log_index].term != request.prev_log_term:
                # Conflicting entry - delete and all following
                self.log = self.log[:request.prev_log_index]
                return AppendEntriesResponse(
                    term=self.current_term,
                    success=False
                )
        
        # Append new entries
        for entry in request.entries:
            if entry.index < len(self.log):
                if self.log[entry.index].term != entry.term:
                    self.log = self.log[:entry.index]
                    self.log.append(entry)
            else:
                self.log.append(entry)
        
        # Update commit index
        if request.leader_commit > self.commit_index:
            self.commit_index = min(
                request.leader_commit,
                len(self.log) - 1 if self.log else 0
            )
        
        return AppendEntriesResponse(
            term=self.current_term,
            success=True
        )
    
    # ==========================================================================
    # Client API
    # ==========================================================================
    
    def is_leader(self) -> bool:
        """Check if this node is the leader."""
        return self.state == NodeState.LEADER
    
    def get_leader_info(self) -> dict:
        """Get information about current leadership."""
        return {
            "node_id": self.node_id,
            "state": self.state.value,
            "term": self.current_term,
            "is_leader": self.is_leader()
        }

# Leader Election with ZooKeeper

from kazoo.client import KazooClient
from kazoo.recipe.election import Election
from kazoo.recipe.watchers import DataWatch
import threading
import logging

logger = logging.getLogger(__name__)


class ZooKeeperLeaderElection:
    """
    Leader election using ZooKeeper.
    
    Uses ephemeral sequential nodes for fair election.
    """
    
    def __init__(
        self,
        zk_hosts: str,
        election_path: str,
        node_id: str,
        on_leader: callable,
        on_follower: callable
    ):
        self.zk = KazooClient(hosts=zk_hosts)
        self.election_path = election_path
        self.node_id = node_id
        self.on_leader = on_leader
        self.on_follower = on_follower
        
        self.is_leader = False
        self.my_node: str = None
        self.running = False
    
    def start(self):
        """Start the leader election."""
        self.zk.start()
        self.running = True
        
        # Ensure election path exists
        self.zk.ensure_path(self.election_path)
        
        # Create ephemeral sequential node
        self.my_node = self.zk.create(
            f"{self.election_path}/candidate-",
            value=self.node_id.encode(),
            ephemeral=True,
            sequence=True
        )
        
        logger.info(f"Created election node: {self.my_node}")
        
        # Check if we're the leader
        self._check_leadership()
    
    def stop(self):
        """Stop the leader election."""
        self.running = False
        if self.my_node:
            try:
                self.zk.delete(self.my_node)
            except Exception:
                pass
        self.zk.stop()
    
    def _check_leadership(self):
        """Check if we're the leader."""
        # Get all candidates
        candidates = self.zk.get_children(self.election_path)
        candidates.sort()  # Sort by sequence number
        
        my_name = self.my_node.split("/")[-1]
        my_index = candidates.index(my_name)
        
        if my_index == 0:
            # We're first - we're the leader!
            if not self.is_leader:
                self.is_leader = True
                logger.info(f"Node {self.node_id} became leader")
                self.on_leader()
        else:
            # Watch the node before us
            self.is_leader = False
            watch_node = f"{self.election_path}/{candidates[my_index - 1]}"
            
            @self.zk.DataWatch(watch_node)
            def watch_predecessor(data, stat, event):
                if event and event.type == "DELETED":
                    # Predecessor died - check leadership again
                    self._check_leadership()
                    return False  # Stop watching
                return True  # Keep watching
            
            self.on_follower()
    
    def get_leader(self) -> str:
        """Get current leader's node ID."""
        candidates = self.zk.get_children(self.election_path)
        candidates.sort()
        
        if candidates:
            leader_path = f"{self.election_path}/{candidates[0]}"
            data, _ = self.zk.get(leader_path)
            return data.decode()
        
        return None


# Using Kazoo's built-in Election recipe
class KazooLeaderElection:
    """
    Simpler leader election using Kazoo's Election recipe.
    """
    
    def __init__(
        self,
        zk_hosts: str,
        election_path: str,
        node_id: str
    ):
        self.zk = KazooClient(hosts=zk_hosts)
        self.election = None
        self.election_path = election_path
        self.node_id = node_id
        self.leader_event = threading.Event()
    
    def start(self):
        """Start and run for election."""
        self.zk.start()
        self.election = Election(self.zk, self.election_path, self.node_id)
    
    def run_for_leader(self, leader_func: callable):
        """
        Run for leader - blocks until elected.
        
        leader_func is called when we become leader.
        When leader_func returns (or raises), leadership is released.
        """
        self.election.run(leader_func)
    
    def contenders(self) -> list:
        """Get list of election contenders."""
        return self.election.contenders()
    
    def stop(self):
        """Stop the election."""
        self.election.cancel()
        self.zk.stop()


# Example usage
def leader_task():
    """Task to run while we're the leader."""
    logger.info("I am the leader! Starting leader duties...")
    
    try:
        while True:
            # Do leader work
            process_as_leader()
            time.sleep(1)
    except Exception as e:
        logger.error(f"Leader task failed: {e}")
        raise  # Will release leadership


def main():
    election = KazooLeaderElection(
        zk_hosts="localhost:2181",
        election_path="/my-service/leader",
        node_id="node-1"
    )
    
    election.start()
    
    # This blocks until we're elected, then runs leader_task
    # When leader_task returns/fails, we release leadership
    election.run_for_leader(leader_task)

# Distributed Locks with ZooKeeper

from kazoo.client import KazooClient
from kazoo.recipe.lock import Lock
from kazoo.exceptions import LockTimeout
from contextlib import contextmanager
import time


class DistributedLock:
    """
    Distributed lock using ZooKeeper.
    
    Features:
    - Automatic release when the session expires (ephemeral nodes)
    - Fair ordering (sequential nodes)
    - Fencing token support
    """
    
    def __init__(self, zk_hosts: str, lock_path: str, identifier: str):
        self.zk = KazooClient(hosts=zk_hosts)
        self.lock_path = lock_path
        self.identifier = identifier
        self.lock = None
        self.fencing_token: int = None
    
    def start(self):
        """Connect to ZooKeeper."""
        self.zk.start()
        self.lock = Lock(self.zk, self.lock_path, self.identifier)
    
    def stop(self):
        """Disconnect from ZooKeeper."""
        self.zk.stop()
    
    def acquire(self, timeout: float = None, blocking: bool = True) -> bool:
        """
        Acquire the lock.
        
        Returns True if acquired, False if the timeout expires or a
        non-blocking attempt fails (kazoo raises LockTimeout on timeout).
        """
        try:
            acquired = self.lock.acquire(blocking=blocking, timeout=timeout)
        except LockTimeout:
            acquired = False
        
        if acquired:
            # Get fencing token from our node's sequence number
            self.fencing_token = self._get_fencing_token()
        
        return acquired
    
    def release(self):
        """Release the lock."""
        self.lock.release()
        self.fencing_token = None
    
    def get_fencing_token(self) -> int:
        """Get the fencing token for this lock acquisition."""
        return self.fencing_token
    
    def _get_fencing_token(self) -> int:
        """Extract fencing token from ZooKeeper sequential node."""
        # The lock creates nodes like /lock/lock-0000000001
        # Extract the sequence number as fencing token
        if hasattr(self.lock, 'node') and self.lock.node:
            node_name = self.lock.node.split("/")[-1]
            # Extract sequence number (last 10 digits)
            return int(node_name[-10:])
        return 0
    
    @contextmanager
    def acquired(self, timeout: float = None):
        """Context manager for lock acquisition."""
        if self.acquire(timeout=timeout):
            try:
                yield self.fencing_token
            finally:
                self.release()
        else:
            raise LockAcquisitionError("Failed to acquire lock")


# Example: Protected critical section
async def critical_section_example():
    """Example of using distributed lock with fencing."""
    lock = DistributedLock(
        zk_hosts="localhost:2181",
        lock_path="/locks/critical-resource",
        identifier="worker-1"
    )
    lock.start()
    
    try:
        with lock.acquired(timeout=10) as fencing_token:
            print(f"Acquired lock with token: {fencing_token}")
            
            # Use fencing token when accessing resources
            # Store the token with the write; reject anything older
            await database.execute(
                "UPDATE resource SET value = $1, fencing_token = $2 "
                "WHERE id = $3 AND fencing_token <= $2",
                new_value,
                fencing_token,
                resource_id
            )
            
    except LockAcquisitionError:
        print("Could not acquire lock")
    
    finally:
        lock.stop()

# Leader Election and Locks with etcd

import etcd3
from typing import Optional, Callable
import threading
import time


class EtcdLeaderElection:
    """
    Leader election using etcd.
    
    Uses etcd's lease and election mechanisms.
    """
    
    def __init__(
        self,
        etcd_host: str,
        etcd_port: int,
        election_name: str,
        node_id: str,
        lease_ttl: int = 10
    ):
        self.client = etcd3.client(host=etcd_host, port=etcd_port)
        self.election_name = election_name
        self.node_id = node_id
        self.lease_ttl = lease_ttl
        
        self.lease = None
        self.election = None
        self.is_leader = False
    
    def campaign(self, on_elected: Callable):
        """
        Campaign for leadership.
        
        Blocks until elected, then calls on_elected.
        """
        # Create lease
        self.lease = self.client.lease(self.lease_ttl)
        
        # Create election. Pseudocode: python-etcd3 has no election helper
        # (it has leases and locks). etcd's Election API (campaign, resign,
        # leader) is in the Go and Java clients (e.g. Go concurrency.Election);
        # in Python you would build the same thing from a lease plus a
        # compare-and-put transaction on a key.
        self.election = self.client.election(self.election_name)
        
        # Campaign (blocks until elected)
        self.election.campaign(self.node_id.encode(), lease=self.lease)
        
        self.is_leader = True
        
        # Start lease keepalive in background
        keepalive_thread = threading.Thread(target=self._keepalive_loop, daemon=True)
        keepalive_thread.start()
        
        # Run leader function
        try:
            on_elected()
        finally:
            self.resign()
    
    def _keepalive_loop(self):
        """Keep lease alive while we're leader."""
        while self.is_leader:
            try:
                self.lease.refresh()
                time.sleep(self.lease_ttl / 3)
            except Exception as e:
                print(f"Keepalive failed: {e}")
                self.is_leader = False
                break
    
    def resign(self):
        """Resign from leadership."""
        self.is_leader = False
        if self.election:
            self.election.resign()
        if self.lease:
            self.lease.revoke()
    
    def get_leader(self) -> Optional[str]:
        """Get current leader."""
        if self.election:
            leader = self.election.leader()
            if leader:
                return leader.value.decode()
        return None


class EtcdDistributedLock:
    """
    Distributed lock using etcd.
    """
    
    def __init__(
        self,
        etcd_host: str,
        etcd_port: int,
        lock_name: str,
        ttl: int = 30
    ):
        self.client = etcd3.client(host=etcd_host, port=etcd_port)
        self.lock_name = lock_name
        self.ttl = ttl
        self.lock = None
    
    def acquire(self, timeout: float = None) -> bool:
        """Acquire the lock."""
        self.lock = self.client.lock(self.lock_name, ttl=self.ttl)
        return self.lock.acquire(timeout=timeout)
    
    def release(self):
        """Release the lock."""
        if self.lock:
            self.lock.release()
            self.lock = None
    
    def refresh(self):
        """Refresh the lock's TTL."""
        if self.lock:
            self.lock.refresh()
    
    @property
    def is_acquired(self) -> bool:
        """Check if lock is acquired."""
        return self.lock is not None and self.lock.is_acquired()


# Example: Job scheduler with leader election
class DistributedJobScheduler:
    """
    Job scheduler that uses leader election.
    
    Only the leader schedules jobs.
    """
    
    def __init__(self, etcd_host: str, node_id: str):
        self.election = EtcdLeaderElection(
            etcd_host=etcd_host,
            etcd_port=2379,
            election_name="job-scheduler",
            node_id=node_id
        )
        self.running = False
    
    def start(self):
        """Start the scheduler."""
        self.running = True
        
        # Campaign for leadership (blocks until elected)
        self.election.campaign(self._run_as_leader)
    
    def _run_as_leader(self):
        """Run scheduler as leader."""
        print(f"Became leader, starting job scheduling...")
        
        while self.running and self.election.is_leader:
            # Schedule jobs
            self._schedule_pending_jobs()
            time.sleep(1)
        
        print("No longer leader, stopping scheduling")
    
    def _schedule_pending_jobs(self):
        """Schedule pending jobs (leader only)."""
        # Implementation would fetch pending jobs and assign to workers
        pass
    
    def stop(self):
        """Stop the scheduler."""
        self.running = False
        self.election.resign()

# Lease-Based Leadership

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Optional
import asyncio


@dataclass
class Lease:
    """A time-limited lease for leadership."""
    holder: str
    token: int
    expires_at: datetime
    
    def is_valid(self) -> bool:
        return datetime.now(timezone.utc) < self.expires_at
    
    def remaining_time(self) -> timedelta:
        return self.expires_at - datetime.now(timezone.utc)


class LeaseBasedLeader:
    """
    Leader election using time-based leases.
    
    Leader must renew lease before expiration.
    If lease expires, another node can claim leadership.
    """
    
    def __init__(
        self,
        node_id: str,
        lease_store,  # Database or coordination service
        lease_duration: timedelta = timedelta(seconds=30),
        renewal_interval: timedelta = timedelta(seconds=10)
    ):
        self.node_id = node_id
        self.store = lease_store
        self.lease_duration = lease_duration
        self.renewal_interval = renewal_interval
        
        self.current_lease: Optional[Lease] = None
        self.is_leader = False
        self.running = False
    
    async def start(self):
        """Start the leader election process."""
        self.running = True
        asyncio.create_task(self._election_loop())
    
    async def stop(self):
        """Stop and release leadership."""
        self.running = False
        if self.is_leader:
            await self._release_lease()
    
    async def _election_loop(self):
        """Main election loop."""
        while self.running:
            try:
                if self.is_leader:
                    await self._renew_lease()
                else:
                    await self._try_acquire_lease()
            except Exception as e:
                logger.error(f"Election loop error: {e}")
                self.is_leader = False
            
            await asyncio.sleep(self.renewal_interval.total_seconds())
    
    async def _try_acquire_lease(self):
        """Try to acquire leadership lease."""
        # Atomic compare-and-set in database
        result = await self.store.try_acquire_lease(
            lease_key="leader",
            holder=self.node_id,
            duration=self.lease_duration
        )
        
        if result.acquired:
            self.current_lease = Lease(
                holder=self.node_id,
                token=result.token,
                expires_at=result.expires_at
            )
            self.is_leader = True
            logger.info(f"Acquired leadership with token {result.token}")
    
    async def _renew_lease(self):
        """Renew the leadership lease."""
        if not self.current_lease or not self.current_lease.is_valid():
            self.is_leader = False
            return
        
        result = await self.store.renew_lease(
            lease_key="leader",
            holder=self.node_id,
            current_token=self.current_lease.token,
            duration=self.lease_duration
        )
        
        if result.renewed:
            self.current_lease = Lease(
                holder=self.node_id,
                token=result.token,
                expires_at=result.expires_at
            )
        else:
            # Lost leadership
            self.is_leader = False
            self.current_lease = None
            logger.warning("Lost leadership - lease not renewed")
    
    async def _release_lease(self):
        """Explicitly release the lease."""
        if self.current_lease:
            await self.store.release_lease(
                lease_key="leader",
                holder=self.node_id,
                token=self.current_lease.token
            )
            self.current_lease = None
            self.is_leader = False
    
    def get_fencing_token(self) -> Optional[int]:
        """Get fencing token for leader operations."""
        if self.current_lease and self.is_leader:
            return self.current_lease.token
        return None


# Database-backed lease store
class PostgresLeaseStore:
    """
    Lease store backed by PostgreSQL.
    
    Uses atomic conditional upserts (no advisory locks needed).
    """
    
    def __init__(self, db_pool):
        self.db = db_pool
    
    async def try_acquire_lease(
        self,
        lease_key: str,
        holder: str,
        duration: timedelta
    ) -> dict:
        """Try to acquire a lease atomically."""
        expires_at = datetime.now(timezone.utc) + duration
        
        # Atomic: Insert if not exists, or update if expired
        result = await self.db.fetchrow("""
            INSERT INTO leases (lease_key, holder, token, expires_at)
            VALUES ($1, $2, 1, $3)
            ON CONFLICT (lease_key) DO UPDATE
            SET holder = $2,
                token = leases.token + 1,
                expires_at = $3
            WHERE leases.expires_at < NOW()
            RETURNING token, expires_at
        """, lease_key, holder, expires_at)
        
        if result:
            return {
                "acquired": True,
                "token": result["token"],
                "expires_at": result["expires_at"]
            }
        
        return {"acquired": False}
    
    async def renew_lease(
        self,
        lease_key: str,
        holder: str,
        current_token: int,
        duration: timedelta
    ) -> dict:
        """Renew an existing lease."""
        expires_at = datetime.now(timezone.utc) + duration
        
        # Only renew if we still hold the lease with the same token
        result = await self.db.fetchrow("""
            UPDATE leases
            SET expires_at = $4,
                token = token + 1
            WHERE lease_key = $1
              AND holder = $2
              AND token = $3
            RETURNING token, expires_at
        """, lease_key, holder, current_token, expires_at)
        
        if result:
            return {
                "renewed": True,
                "token": result["token"],
                "expires_at": result["expires_at"]
            }
        
        return {"renewed": False}
    
    async def release_lease(
        self,
        lease_key: str,
        holder: str,
        token: int
    ):
        """Release a lease."""
        await self.db.execute("""
            DELETE FROM leases
            WHERE lease_key = $1
              AND holder = $2
              AND token = $3
        """, lease_key, holder, token)

# Quorum-Based Leader Election

class QuorumLeaderElection:
    """
    Leader election requiring quorum agreement.
    
    Prevents split-brain by requiring majority.
    """
    
    def __init__(
        self,
        node_id: str,
        all_nodes: list[str],
        rpc_client
    ):
        self.node_id = node_id
        self.all_nodes = all_nodes
        self.rpc = rpc_client
        self.quorum_size = len(all_nodes) // 2 + 1
        
        self.current_term = 0
        self.voted_for: Optional[str] = None
        self.is_leader = False
    
    async def try_become_leader(self) -> bool:
        """
        Attempt to become leader.
        
        Returns True if elected with quorum support.
        """
        self.current_term += 1
        self.voted_for = self.node_id
        
        votes = 1  # Vote for self
        
        # Request votes from all other nodes
        for node in self.all_nodes:
            if node == self.node_id:
                continue
            
            try:
                response = await self.rpc.request_vote(
                    node,
                    term=self.current_term,
                    candidate=self.node_id
                )
                
                if response.vote_granted:
                    votes += 1
                
            except Exception as e:
                logger.warning(f"Failed to get vote from {node}: {e}")
        
        # Check if we have quorum
        if votes >= self.quorum_size:
            self.is_leader = True
            logger.info(f"Elected leader with {votes}/{len(self.all_nodes)} votes")
            return True
        
        logger.info(f"Election failed: {votes}/{self.quorum_size} votes needed")
        return False
    
    async def verify_leadership(self) -> bool:
        """
        Verify we still have quorum support.
        
        Call periodically to ensure we're still the valid leader.
        """
        if not self.is_leader:
            return False
        
        confirmations = 1  # Self
        
        for node in self.all_nodes:
            if node == self.node_id:
                continue
            
            try:
                response = await self.rpc.heartbeat(
                    node,
                    term=self.current_term,
                    leader=self.node_id
                )
                
                if response.acknowledged:
                    confirmations += 1
                    
            except Exception:
                pass  # Node unreachable
        
        if confirmations < self.quorum_size:
            # Lost quorum - step down
            self.is_leader = False
            logger.warning("Lost quorum - stepping down")
            return False
        
        return True