use deadline::deadline;
use once_cell::sync::Lazy;
use parking_lot::Mutex;

mod common;
use std::{collections::HashSet, net::SocketAddr, sync::Arc, time::Duration};

use pea2pea::{
    protocols::{OnDisconnect, Reading},
    Pea2Pea,
};

static DISCONNECT_TRIGGERED: Lazy<Arc<Mutex<HashSet<String>>>> = Lazy::new(Default::default);

impl OnDisconnect for common::TestNode {
    async fn on_disconnect(&self, _addr: SocketAddr) {
        DISCONNECT_TRIGGERED
            .lock()
            .insert(self.node().name().to_owned());
    }
}

#[tokio::test]
async fn connector_side_disconnect() {
    let connector = crate::test_node!("connector0");
    connector.enable_on_disconnect().await;

    let connectee = crate::test_node!("connectee0");
    connectee.enable_reading().await;
    connectee.enable_on_disconnect().await;
    let connectee_addr = connectee.node().toggle_listener().await.unwrap().unwrap();

    connector.node().connect(connectee_addr).await.unwrap();

    let connectee_clone = connectee.clone();
    deadline!(Duration::from_secs(1), move || connectee_clone
        .node()
        .num_connected()
        == 1);

    let disconnects = DISCONNECT_TRIGGERED.lock().clone();
    assert!(
        !disconnects.contains(connector.node().name())
            && !disconnects.contains(connectee.node().name())
    );

    connector.node().disconnect(connectee_addr).await;

    deadline!(Duration::from_secs(1), move || {
        let disconnects = DISCONNECT_TRIGGERED.lock().clone();
        disconnects.contains(connector.node().name())
            && disconnects.contains(connectee.node().name())
    });
}

#[tokio::test]
async fn connectee_side_disconnect() {
    let connector = crate::test_node!("connector1");
    connector.enable_reading().await;
    connector.enable_on_disconnect().await;

    let connectee = crate::test_node!("connectee1");
    connectee.enable_on_disconnect().await;
    let connectee_addr = connectee.node().toggle_listener().await.unwrap().unwrap();

    connector.node().connect(connectee_addr).await.unwrap();

    let connectee_clone = connectee.clone();
    deadline!(Duration::from_secs(1), move || connectee_clone
        .node()
        .num_connected()
        == 1);

    let disconnects = DISCONNECT_TRIGGERED.lock().clone();
    assert!(
        !disconnects.contains(connector.node().name())
            && !disconnects.contains(connectee.node().name())
    );

    let connector_addr = connectee.node().connected_addrs()[0];
    connectee.node().disconnect(connector_addr).await;

    deadline!(Duration::from_secs(1), move || {
        let disconnects = DISCONNECT_TRIGGERED.lock().clone();
        disconnects.contains(connector.node().name())
            && disconnects.contains(connectee.node().name())
    });
}