grapevine 1.1.0

//! Peer management types.

use std::net::SocketAddr;
use std::time::{Duration, Instant};

use tokio::sync::mpsc::Sender;
use tokio::sync::mpsc::error::TrySendError;
use tracing::debug;

use crate::{Error, Message, Result};

/// Age bonus divisor for health score calculation (seconds).
const AGE_BONUS_DIVISOR: f64 = 300.0;

/// Maximum age bonus for health score.
const MAX_AGE_BONUS: f64 = 0.2;

/// Penalty per consecutive failure for health score.
const CONSECUTIVE_FAILURE_PENALTY: f64 = 0.1;

/// Maximum consecutive failure penalty for health score.
const MAX_CONSECUTIVE_PENALTY: f64 = 0.5;

/// Maximum consecutive failures before peer disconnection.
const MAX_CONSECUTIVE_FAILURES: u64 = 5;

/// State of a peer connection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PeerState {
    /// Connection is being established
    Connecting,
    /// Connection is established and healthy
    Connected,
    /// Peer is connected but unresponsive
    Stale,
    /// Peer is disconnected
    Disconnected,
}

/// Information about a single live connection.
#[derive(Debug, Clone)]
pub struct PeerInfo {
    /// Connection address of the peer
    pub addr: SocketAddr,

    /// Current state
    pub state: PeerState,

    /// Last time we received a message from this peer
    pub last_seen: Instant,

    /// When the connection was established
    pub connected_at: Instant,

    /// Number of messages received from this peer
    pub messages_received: u64,

    /// Number of messages sent to this peer
    pub messages_sent: u64,

    /// Number of failed send attempts
    pub message_failures: u64,

    /// Consecutive failures (reset on success)
    pub consecutive_failures: u64,
}

impl PeerInfo {
    /// Create new peer info for a connection at `addr`.
    pub fn new(addr: SocketAddr) -> Self {
        let now = Instant::now();
        Self {
            addr,
            state: PeerState::Connecting,
            last_seen: now,
            connected_at: now,
            messages_received: 0,
            messages_sent: 0,
            message_failures: 0,
            consecutive_failures: 0,
        }
    }

    /// Mark peer as connected.
    pub fn mark_connected(&mut self) {
        self.state = PeerState::Connected;
    }

    /// Mark peer as stale.
    pub fn mark_stale(&mut self) {
        self.state = PeerState::Stale;
    }

    /// Mark peer as disconnected.
    pub fn mark_disconnected(&mut self) {
        self.state = PeerState::Disconnected;
    }

    /// Update last seen timestamp.
    pub fn update_last_seen(&mut self) {
        self.last_seen = Instant::now();
        if matches!(self.state, PeerState::Connecting | PeerState::Stale) {
            self.state = PeerState::Connected;
        }
    }

    /// Check if peer is stale (hasn't been seen recently).
    pub fn is_stale(&self, timeout: Duration) -> bool {
        self.last_seen.elapsed() > timeout
    }

    /// Increment received message counter.
    pub fn increment_received(&mut self) {
        self.messages_received = self.messages_received.saturating_add(1);
        self.update_last_seen();
    }

    /// Increment sent message counter.
    pub fn increment_sent(&mut self) {
        self.messages_sent = self.messages_sent.saturating_add(1);
        self.consecutive_failures = 0;
    }

    /// Record a failure in sending a message.
    pub fn record_failure(&mut self) {
        self.message_failures = self.message_failures.saturating_add(1);
        self.consecutive_failures = self.consecutive_failures.saturating_add(1);
    }

    /// Calculate health score (0.0 = poor, 1.0 = excellent).
    pub fn health_score(&self) -> f64 {
        let total_attempts = self.messages_sent.saturating_add(self.message_failures);
        if total_attempts == 0 {
            return 1.0;
        }

        let sent = f64::from(u32::try_from(self.messages_sent).unwrap_or(u32::MAX));
        let total = f64::from(u32::try_from(total_attempts).unwrap_or(u32::MAX));
        let success_rate = sent / total;

        let age_seconds =
            f64::from(u32::try_from(self.connected_at.elapsed().as_secs()).unwrap_or(u32::MAX));
        let age_bonus = (age_seconds / AGE_BONUS_DIVISOR).min(MAX_AGE_BONUS);

        let consecutive_penalty =
            (f64::from(u32::try_from(self.consecutive_failures).unwrap_or(u32::MAX))
                * CONSECUTIVE_FAILURE_PENALTY)
                .min(MAX_CONSECUTIVE_PENALTY);

        (success_rate + age_bonus - consecutive_penalty).clamp(0.0, 1.0)
    }

    /// Check if peer should be disconnected due to failures.
    pub fn should_disconnect(&self) -> bool {
        self.consecutive_failures >= MAX_CONSECUTIVE_FAILURES
    }
}

/// Peer handle for network operations.
#[derive(Debug)]
pub struct Peer {
    /// Peer information
    pub info: PeerInfo,
    /// Bounded channel to this peer's writer task.
    sender: Sender<Message>,
}

impl Peer {
    /// Create a new peer for the connection at `addr`.
    pub fn new(addr: SocketAddr, sender: Sender<Message>) -> Self {
        Self {
            info: PeerInfo::new(addr),
            sender,
        }
    }

    /// Queue a message for this peer's writer task.
    pub fn send(&mut self, message: Message) -> Result<()> {
        match self.sender.try_send(message) {
            Ok(()) => {
                self.info.increment_sent();
                Ok(())
            }
            Err(TrySendError::Full(_)) => {
                debug!(
                    "Write channel full for {}, dropping message",
                    self.info.addr
                );
                Ok(())
            }
            Err(TrySendError::Closed(_)) => {
                self.info.record_failure();
                Err(Error::Channel("peer write channel closed".to_string()))
            }
        }
    }

    /// Get the peer's connection address.
    pub fn addr(&self) -> SocketAddr {
        self.info.addr
    }

    /// Get peer state.
    pub fn state(&self) -> PeerState {
        self.info.state
    }
}

#[cfg(test)]
mod tests {
    use bytes::Bytes;

    use super::*;
    use crate::Payload;

    fn message(seq: u64) -> Message {
        let addr = "127.0.0.1:8000".parse().unwrap();
        Message::new(
            addr,
            seq,
            Payload::Application(Bytes::from(format!("m{seq}"))),
        )
    }

    #[test]
    fn update_last_seen_drives_state_machine() {
        let addr = "127.0.0.1:8000".parse().unwrap();
        let mut info = PeerInfo::new(addr);
        assert_eq!(info.state, PeerState::Connecting);

        info.update_last_seen();
        assert_eq!(info.state, PeerState::Connected);

        info.mark_stale();
        assert_eq!(info.state, PeerState::Stale);
        info.update_last_seen();
        assert_eq!(info.state, PeerState::Connected);

        info.mark_disconnected();
        info.update_last_seen();
        assert_eq!(info.state, PeerState::Disconnected);
    }

    #[test]
    fn peer_info_saturating_counters() {
        let addr = "127.0.0.1:8000".parse().unwrap();
        let mut info = PeerInfo::new(addr);

        info.messages_received = u64::MAX;
        info.increment_received();
        assert_eq!(info.messages_received, u64::MAX);

        info.messages_sent = u64::MAX;
        info.increment_sent();
        assert_eq!(info.messages_sent, u64::MAX);
    }

    #[test]
    fn peer_send_success() {
        let addr = "127.0.0.1:8000".parse().unwrap();
        let (tx, mut rx) = tokio::sync::mpsc::channel(16);

        let mut peer = Peer::new(addr, tx);
        let msg = message(1);

        assert_eq!(peer.info.messages_sent, 0);
        assert!(peer.send(msg.clone()).is_ok());
        assert_eq!(peer.info.messages_sent, 1);

        let received = rx.try_recv().expect("message was queued");
        assert_eq!(received.id, msg.id);
    }

    #[test]
    fn peer_send_failure_channel_closed() {
        let addr = "127.0.0.1:8000".parse().unwrap();
        let (tx, rx) = tokio::sync::mpsc::channel::<Message>(16);

        // Drop receiver to close channel
        drop(rx);

        let mut peer = Peer::new(addr, tx);

        let result = peer.send(message(1));
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::Channel(_)));
    }

    #[test]
    fn peer_multiple_sends() {
        let addr = "127.0.0.1:8000".parse().unwrap();
        let (tx, mut rx) = tokio::sync::mpsc::channel(16);

        let mut peer = Peer::new(addr, tx);

        for i in 1..=5 {
            assert!(peer.send(message(i)).is_ok());
            assert_eq!(peer.info.messages_sent, i);

            let received = rx.try_recv().unwrap();
            assert_eq!(received.id.sequence, i);
        }
    }

    #[test]
    fn send_drops_newest_when_channel_full() {
        let addr = "127.0.0.1:8000".parse().unwrap();
        // Capacity 1, receiver never drains: the second send overflows.
        let (tx, _rx) = tokio::sync::mpsc::channel::<Message>(1);

        let mut peer = Peer::new(addr, tx);

        assert!(peer.send(message(1)).is_ok());
        assert_eq!(peer.info.messages_sent, 1);

        // Full channel: dropped, still Ok, and not counted as a send.
        assert!(peer.send(message(2)).is_ok());
        assert_eq!(
            peer.info.messages_sent, 1,
            "a dropped frame is not counted as sent"
        );
        assert_eq!(
            peer.info.consecutive_failures, 0,
            "backpressure is not a peer failure"
        );
    }
}