//! Ban list for misbehaving peers.
//!
//! Tracks failure counts per peer. After exceeding a
//! threshold, the peer is temporarily banned. Bans
//! expire automatically after a configurable duration.

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

/// Default number of failures before banning a peer.
const DEFAULT_BAN_THRESHOLD: u32 = 3;

/// Default ban duration (3 hours, matching gonode).
const DEFAULT_BAN_DURATION: Duration = Duration::from_secs(3 * 3600);

/// Tracks failures and bans for peers.
pub struct BanList {
    /// Failure counts per address.
    failures: HashMap<SocketAddr, FailureEntry>,
    /// Active bans: address → expiry time.
    bans: HashMap<SocketAddr, Instant>,
    /// Number of failures before a ban is applied.
    threshold: u32,
    /// How long a ban lasts.
    ban_duration: Duration,
}

struct FailureEntry {
    count: u32,
    last_failure: Instant,
}

impl BanList {
    pub fn new() -> Self {
        Self {
            failures: HashMap::new(),
            bans: HashMap::new(),
            threshold: DEFAULT_BAN_THRESHOLD,
            ban_duration: DEFAULT_BAN_DURATION,
        }
    }

    /// Set the failure threshold before banning.
    pub fn set_threshold(&mut self, threshold: u32) {
        self.threshold = threshold;
    }

    /// Set the ban duration.
    pub fn set_ban_duration(&mut self, duration: Duration) {
        self.ban_duration = duration;
    }

    /// Check if a peer is currently banned.
    pub fn is_banned(&self, addr: &SocketAddr) -> bool {
        if let Some(expiry) = self.bans.get(addr) {
            Instant::now() < *expiry
        } else {
            false
        }
    }

    /// Record a failure for a peer. Returns true if the
    /// peer was just banned (crossed the threshold).
    pub fn record_failure(&mut self, addr: SocketAddr) -> bool {
        let entry = self.failures.entry(addr).or_insert(FailureEntry {
            count: 0,
            last_failure: Instant::now(),
        });
        entry.count += 1;
        entry.last_failure = Instant::now();

        if entry.count >= self.threshold {
            let expiry = Instant::now() + self.ban_duration;
            self.bans.insert(addr, expiry);
            self.failures.remove(&addr);
            log::info!(
                "Banned peer {addr} for {}s",
                self.ban_duration.as_secs()
            );
            true
        } else {
            false
        }
    }

    /// Clear failure count for a peer (e.g. after a
    /// successful interaction).
    pub fn record_success(&mut self, addr: &SocketAddr) {
        self.failures.remove(addr);
    }

    /// Remove expired bans and stale failure entries.
    /// Call periodically from the event loop.
    pub fn cleanup(&mut self) {
        let now = Instant::now();
        self.bans.retain(|_, expiry| now < *expiry);

        // Clear failure entries older than ban_duration
        // (stale failures shouldn't accumulate forever)
        self.failures
            .retain(|_, e| e.last_failure.elapsed() < self.ban_duration);
    }

    /// Number of currently active bans.
    pub fn ban_count(&self) -> usize {
        self.bans
            .iter()
            .filter(|(_, e)| Instant::now() < **e)
            .count()
    }

    /// Number of peers with recorded failures.
    pub fn failure_count(&self) -> usize {
        self.failures.len()
    }
}

impl Default for BanList {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn addr(port: u16) -> SocketAddr {
        SocketAddr::from(([127, 0, 0, 1], port))
    }

    #[test]
    fn not_banned_initially() {
        let bl = BanList::new();
        assert!(!bl.is_banned(&addr(1000)));
    }

    #[test]
    fn ban_after_threshold() {
        let mut bl = BanList::new();
        let a = addr(1000);
        assert!(!bl.record_failure(a));
        assert!(!bl.record_failure(a));
        assert!(bl.record_failure(a)); // 3rd failure → banned
        assert!(bl.is_banned(&a));
    }

    #[test]
    fn success_clears_failures() {
        let mut bl = BanList::new();
        let a = addr(1000);
        bl.record_failure(a);
        bl.record_failure(a);
        bl.record_success(&a);
        // Failures cleared, next failure starts over
        assert!(!bl.record_failure(a));
        assert!(!bl.is_banned(&a));
    }

    #[test]
    fn ban_expires() {
        let mut bl = BanList::new();
        bl.set_ban_duration(Duration::from_millis(1));
        let a = addr(1000);
        bl.record_failure(a);
        bl.record_failure(a);
        bl.record_failure(a);
        assert!(bl.is_banned(&a));
        std::thread::sleep(Duration::from_millis(5));
        assert!(!bl.is_banned(&a));
    }

    #[test]
    fn cleanup_removes_expired() {
        let mut bl = BanList::new();
        bl.set_ban_duration(Duration::from_millis(1));
        let a = addr(1000);
        bl.record_failure(a);
        bl.record_failure(a);
        bl.record_failure(a);
        std::thread::sleep(Duration::from_millis(5));
        bl.cleanup();
        assert_eq!(bl.ban_count(), 0);
    }

    #[test]
    fn custom_threshold() {
        let mut bl = BanList::new();
        bl.set_threshold(1);
        let a = addr(1000);
        assert!(bl.record_failure(a)); // 1 failure → banned
        assert!(bl.is_banned(&a));
    }

    #[test]
    fn independent_peers() {
        let mut bl = BanList::new();
        let a = addr(1000);
        let b = addr(2000);
        bl.record_failure(a);
        bl.record_failure(a);
        bl.record_failure(a);
        assert!(bl.is_banned(&a));
        assert!(!bl.is_banned(&b));
    }
}