"""
Paxos Group Implementation

Implements Multi-Paxos with optimizations from Spanner:
- Leader leases (skip prepare phase when leader has valid lease)
- Batching and pipelining
- Fast path for common case

Based on:
- "Paxos Made Simple" by Leslie Lamport
- "Paxos Made Live" (Chubby paper)
- Spanner paper optimizations
"""

import threading
import time
from enum import Enum
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Callable
import logging

from spymonk_enterprise.time.hybrid_clock import Timestamp, HybridLogicalClock

logger = logging.getLogger(__name__)


class PaxosState(Enum):
    """State of a Paxos replica"""
    FOLLOWER = "follower"
    CANDIDATE = "candidate"  # Election in progress
    LEADER = "leader"


class ProposeResult(Enum):
    """Result of proposal"""
    SUCCESS = "success"
    FAILED = "failed"
    NOT_LEADER = "not_leader"
    TIMEOUT = "timeout"


@dataclass
class PaxosProposal:
    """
    Single Paxos proposal.

    In Spanner terms, this is a mutation being replicated.
    """
    proposal_number: int           # Monotonically increasing
    value: bytes                   # Serialized mutation
    timestamp: Timestamp           # HLC timestamp
    client_id: Optional[str] = None

    def __lt__(self, other):
        return self.proposal_number < other.proposal_number


@dataclass
class PaxosLogEntry:
    """Entry in the Paxos log"""
    index: int                     # Log position
    proposal: PaxosProposal
    committed: bool = False


@dataclass
class LeaderLease:
    """
    Leader lease for optimization.

    Spanner optimization: Leader can skip prepare phase
    while it holds a valid lease.
    """
    leader_id: str
    start_time: Timestamp
    expiry_time: Timestamp

    def is_valid(self, current_time: Timestamp) -> bool:
        """Check if lease is still valid"""
        return current_time.before(self.expiry_time)


class PaxosGroup:
    """
    Multi-Paxos group for replicating a tablet.

    Features:
    - Multi-Paxos: Elect leader once, then stream accepts
    - Leader leases: Skip prepare phase (Spanner optimization)
    - Batching: Batch multiple proposals
    - State machine replication

    This is the core of Spanner's replication.
    """

    def __init__(
        self,
        group_id: str,
        replica_id: str,
        replicas: List[str],
        clock: HybridLogicalClock,
        state_machine: Optional[Callable] = None,
        lease_duration_ms: int = 10000,  # 10 seconds
    ):
        """
        Initialize Paxos group.

        Args:
            group_id: Unique group identifier
            replica_id: This replica's ID
            replicas: List of all replica IDs
            clock: Hybrid logical clock
            state_machine: Optional callback for applying committed entries
            lease_duration_ms: Leader lease duration in milliseconds
        """
        self.group_id = group_id
        self.replica_id = replica_id
        self.replicas = replicas
        self.clock = clock
        self.state_machine = state_machine
        self.lease_duration_ms = lease_duration_ms

        # Paxos state
        self.state = PaxosState.FOLLOWER
        self.current_leader: Optional[str] = None
        self.current_term = 0

        # Proposer state
        self.next_proposal_number = 0
        self.highest_promised = -1

        # Acceptor state
        self.accepted_proposal: Optional[PaxosProposal] = None
        self.min_proposal = -1

        # Log (replicated state machine log)
        self.log: List[PaxosLogEntry] = []
        self.commit_index = 0
        self.last_applied = 0

        # Leader lease
        self.current_lease: Optional[LeaderLease] = None

        # Network (in-memory for single-node, will use gRPC for distributed)
        self.message_handlers: Dict[str, Callable] = {}

        # Thread safety
        self._lock = threading.RLock()

        # Quorum size
        self.quorum_size = len(replicas) // 2 + 1

        logger.info(f"Initialized Paxos group {group_id}, replica={replica_id}, quorum={self.quorum_size}")

    def propose(self, value: bytes) -> ProposeResult:
        """
        Propose a value to be replicated.

        This is called by the transaction coordinator to replicate
        a mutation across the Paxos group.

        Args:
            value: Serialized mutation

        Returns:
            ProposeResult indicating success/failure
        """
        with self._lock:
            # Must be leader to propose
            if self.state != PaxosState.LEADER:
                logger.warning(f"Propose failed: not leader (state={self.state})")
                return ProposeResult.NOT_LEADER

            # Check if we have valid lease
            if self._has_valid_lease():
                # Fast path: Use lease to skip prepare phase
                return self._propose_with_lease(value)
            else:
                # Slow path: Full Paxos (prepare + accept)
                return self._propose_full_paxos(value)

    def _has_valid_lease(self) -> bool:
        """Check if this leader has a valid lease"""
        if self.current_lease is None:
            return False

        if self.current_lease.leader_id != self.replica_id:
            return False

        current_time = self.clock.now()
        return self.current_lease.is_valid(current_time)

    def _propose_with_lease(self, value: bytes) -> ProposeResult:
        """
        Propose using leader lease (fast path).

        Spanner optimization: Skip prepare phase!
        """
        # Create proposal
        self.next_proposal_number += 1
        proposal = PaxosProposal(
            proposal_number=self.next_proposal_number,
            value=value,
            timestamp=self.clock.now()
        )

        # Send ACCEPT to all replicas (including self)
        accepts = self._send_accept(proposal)

        # Need quorum
        if len(accepts) >= self.quorum_size:
            # Commit!
            self._commit_proposal(proposal)
            logger.debug(f"Proposal {proposal.proposal_number} committed via lease fast path")
            return ProposeResult.SUCCESS
        else:
            logger.warning(f"Proposal {proposal.proposal_number} failed: only {len(accepts)}/{self.quorum_size} accepts")
            return ProposeResult.FAILED

    def _propose_full_paxos(self, value: bytes) -> ProposeResult:
        """
        Full Paxos with prepare + accept phases.

        Used when leader doesn't have a valid lease.
        """
        # Phase 1: Prepare
        self.next_proposal_number += 1
        proposal_number = self.next_proposal_number

        promises = self._send_prepare(proposal_number)

        if len(promises) < self.quorum_size:
            logger.warning(f"Prepare phase failed: only {len(promises)}/{self.quorum_size} promises")
            return ProposeResult.FAILED

        # Check if any promise contained a previously accepted value
        # (This handles the case where we're recovering from a failed leader)
        highest_accepted = None
        for promise in promises:
            if promise.get('accepted_proposal'):
                accepted = promise['accepted_proposal']
                if highest_accepted is None or accepted > highest_accepted:
                    highest_accepted = accepted

        # Use previously accepted value if it exists, otherwise use our value
        proposal_value = highest_accepted.value if highest_accepted else value

        # Phase 2: Accept
        proposal = PaxosProposal(
            proposal_number=proposal_number,
            value=proposal_value,
            timestamp=self.clock.now()
        )

        accepts = self._send_accept(proposal)

        if len(accepts) >= self.quorum_size:
            self._commit_proposal(proposal)

            # Acquire new lease
            self._acquire_lease()

            logger.debug(f"Proposal {proposal_number} committed via full Paxos")
            return ProposeResult.SUCCESS
        else:
            logger.warning(f"Accept phase failed: only {len(accepts)}/{self.quorum_size} accepts")
            return ProposeResult.FAILED

    def _send_prepare(self, proposal_number: int) -> List[dict]:
        """
        Send PREPARE messages to all replicas.

        Returns:
            List of PROMISE responses
        """
        promises = []

        # Send to all replicas (including self)
        for replica_id in self.replicas:
            promise = self._handle_prepare(replica_id, proposal_number)
            if promise:
                promises.append(promise)

        return promises

    def _handle_prepare(self, from_replica: str, proposal_number: int) -> Optional[dict]:
        """
        Handle PREPARE message (Acceptor role).

        Returns:
            PROMISE if proposal_number > min_proposal, None otherwise
        """
        with self._lock:
            if proposal_number > self.min_proposal:
                self.min_proposal = proposal_number

                return {
                    'replica_id': self.replica_id,
                    'proposal_number': proposal_number,
                    'accepted_proposal': self.accepted_proposal
                }
            else:
                return None

    def _send_accept(self, proposal: PaxosProposal) -> List[dict]:
        """
        Send ACCEPT messages to all replicas.

        Returns:
            List of ACCEPTED responses
        """
        accepts = []

        for replica_id in self.replicas:
            accepted = self._handle_accept(replica_id, proposal)
            if accepted:
                accepts.append(accepted)

        return accepts

    def _handle_accept(self, from_replica: str, proposal: PaxosProposal) -> Optional[dict]:
        """
        Handle ACCEPT message (Acceptor role).

        Returns:
            ACCEPTED if proposal_number >= min_proposal, None otherwise
        """
        with self._lock:
            if proposal.proposal_number >= self.min_proposal:
                self.accepted_proposal = proposal
                self.min_proposal = proposal.proposal_number

                return {
                    'replica_id': self.replica_id,
                    'proposal_number': proposal.proposal_number
                }
            else:
                return None

    def _commit_proposal(self, proposal: PaxosProposal):
        """
        Commit a proposal to the log.

        This is called when quorum is reached.
        """
        with self._lock:
            # Add to log
            entry = PaxosLogEntry(
                index=len(self.log),
                proposal=proposal,
                committed=True
            )
            self.log.append(entry)
            self.commit_index = len(self.log) - 1

            # Apply to state machine
            self._apply_committed_entries()

    def _apply_committed_entries(self):
        """Apply committed log entries to state machine"""
        with self._lock:
            while self.last_applied < self.commit_index:
                self.last_applied += 1
                entry = self.log[self.last_applied]

                if self.state_machine:
                    try:
                        self.state_machine(entry.proposal.value)
                    except Exception as e:
                        logger.error(f"State machine application failed: {e}")

    def _acquire_lease(self):
        """
        Acquire leader lease.

        This allows skipping prepare phase for subsequent proposals.
        """
        current_time = self.clock.now()

        # Lease expires after lease_duration_ms
        expiry_ns = current_time.physical + (self.lease_duration_ms * 1_000_000)
        expiry = Timestamp(
            physical=expiry_ns,
            logical=0,
            uncertainty_ns=current_time.uncertainty_ns
        )

        self.current_lease = LeaderLease(
            leader_id=self.replica_id,
            start_time=current_time,
            expiry_time=expiry
        )

        logger.info(f"Acquired leader lease until {expiry}")

    def become_leader(self) -> bool:
        """
        Initiate leader election.

        Returns:
            True if successfully became leader
        """
        with self._lock:
            self.state = PaxosState.CANDIDATE
            self.current_term += 1

            logger.info(f"Starting leader election for term {self.current_term}")

            # Vote for self
            votes = {self.replica_id}

            # Request votes from other replicas
            # (Simplified: in real implementation, send RequestVote RPCs)

            if len(votes) >= self.quorum_size:
                self.state = PaxosState.LEADER
                self.current_leader = self.replica_id

                # Acquire lease
                self._acquire_lease()

                logger.info(f"Became leader for term {self.current_term}")
                return True
            else:
                self.state = PaxosState.FOLLOWER
                return False

    def get_stats(self) -> dict:
        """Get Paxos group statistics"""
        with self._lock:
            return {
                'group_id': self.group_id,
                'replica_id': self.replica_id,
                'state': self.state.value,
                'current_leader': self.current_leader,
                'current_term': self.current_term,
                'log_length': len(self.log),
                'commit_index': self.commit_index,
                'last_applied': self.last_applied,
                'has_lease': self._has_valid_lease(),
                'lease_expiry': self.current_lease.expiry_time if self.current_lease else None
            }
