bevy_connect 0.19.2

use bevy::prelude::Message as BevyMessage;
use bevy::prelude::*;
use bevy_connect::ClientId;
use bevy_connect::prelude::SessionOptions;
use bevy_connect::{
    Message, SessionPlugin,
    channel::{Channel, SessionConfig},
    commands::{SessionConnectCommand, SessionDisconnectCommand},
    events::{MessageReceivedEvent, SessionConnectedEvent, SessionDisconnectedEvent},
    prelude::SessionPromoteToHostCommand,
};
use serde::{Deserialize, Serialize};
use std::{any::type_name, fmt::Debug};
use std::{
    collections::{HashMap, HashSet},
    thread::sleep,
    time::Duration,
};

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
pub struct Msg {
    pub value: i32,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
pub struct Msg2 {
    pub data: Vec<u8>,
}

pub struct TestGroup {
    port: u16,
    host: TestApp,
    clients: Vec<TestApp>,
}

static KEY: &[u8] = &[
    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 2, 23, 24, 25, 26,
    27, 28, 29, 30, 31, 32,
];

impl TestGroup {
    /// Promotes one of the clients to become the host.
    /// Checks at the end that there is only one host and it is the one that
    /// was promoted.
    /// This only promotes the Msg channel and leaves the other like before.
    /// But now the host & clients variables won't tell the truth anymore so
    /// don't continue testing after this change.
    ///
    /// # Panics
    ///
    /// Fails if the assertions fail.
    pub fn test_promote_new_host(&mut self, port: u16) {
        self.update_all();

        info!("Testing host promotion of {}", type_name::<Msg>());
        self.print_all_destinations::<Msg>("Before migration,");

        let list = self.clients_list_minus_host::<Msg>();
        let new_host = *list
            .iter()
            .next()
            .expect("Cannot promote when no clients available");

        assert!(self.host.is_host::<Msg>(), "Was not host to begin with");

        info!("Promoting {new_host} to host");
        self.host.promote_new_host::<Msg>(new_host, port);
        let promoted = self.client_of_uuid::<Msg>(new_host);
        // Client receives promotion, responds with NewHost
        promoted.app.update();
        // Host tells other clients to reconnect & reconnects to new host
        self.host.app.update();
        // All other clients connect to new host
        for c in &mut self.clients {
            if c.uuid::<Msg>() == new_host {
                // swap the host for continuing tests proper
                std::mem::swap(&mut c.app, &mut self.host.app);
                continue;
            }
            c.app.update();
        }

        assert!(self.host.is_host::<Msg>(), "New host was not host");

        self.assert_all_destinations::<Msg>();
        self.assert_host_uuid_reported::<Msg>();
    }

    pub fn connect_new_client(&mut self) {
        let mut c = TestApp::new_client(self.port, false, None);
        self.host.app.update();
        c.app.update();
        self.clients.push(c);
        self.wait_a_while();
    }

    pub fn connect_new_client_comp_enc(&mut self) {
        let mut c = TestApp::new_client(self.port, true, Some(KEY.to_vec()));
        self.host.app.update();
        c.app.update();
        self.clients.push(c);
        self.wait_a_while();
    }

    pub fn disconnect_one_client(&mut self) {
        if let Some(mut c) = self.clients.pop() {
            c.disconnect_one::<Msg>();
            c.disconnect_one::<Msg2>();
            self.host.app.update();
            c.app.update();
        }
        self.wait_a_while();
    }

    /// # Panics
    ///
    /// Panics if assertions are false
    pub fn test_disconnect<M: Message>(&mut self) {
        retry(|| {
            self.update_all();
            for c in &mut self.clients {
                c.disconnect_one::<M>();
                self.host.app.update();
                c.app.update();
            }
            self.host.disconnect_one::<M>();
            self.update_all();
            self.assert_disconnected::<M>()?;
            Ok::<(), String>(())
        })
        .unwrap();
    }

    pub fn wait_a_while(&mut self) {
        for _ in 0..3 {
            self.update_all();
        }
    }

    // Need to check all M types at once otherwise a single update for an app
    // will consume all events and won't be able to be checked again.
    pub fn test_all_clients_joined(&mut self) {
        self.update_all();

        // Verify that the uuids are all assigned and are present in the list
        // obtained by the destinations.
        self.assert_host_is_in_destinations::<Msg>();
        self.assert_clients_are_in_destinations::<Msg>();
        self.assert_uuids_all_unique::<Msg>();
        self.assert_known_host::<Msg>();
        self.assert_all_destinations::<Msg>();

        self.assert_host_is_in_destinations::<Msg2>();
        self.assert_clients_are_in_destinations::<Msg2>();
        self.assert_uuids_all_unique::<Msg2>();
        self.assert_known_host::<Msg2>();
        self.assert_all_destinations::<Msg2>();
    }

    /// This will send and receive the message from all points to all points
    /// and the combination of those between host and client(s).
    pub fn test_message_all_combos<M>(&mut self, m: &M)
    where
        M: Message + Clone + Debug + PartialEq,
    {
        self.wait_a_while();
        info!("Testing broadcast of {}:{:?}", type_name::<M>(), m);
        self.broadcast_receive(&m.clone(), None);
        for i in 0..self.clients.len() {
            self.wait_a_while();
            self.broadcast_receive(&m.clone(), Some(i));
        }
    }

    /// This will send a message only to a specific client picked from the
    /// client list and verify that only that client receives it.
    ///
    /// # Panics
    ///
    /// Fails if the assertions fail.
    pub fn test_message_p2p<M>(&mut self, m: &M)
    where
        M: Message + Clone + Debug + PartialEq,
    {
        self.wait_a_while();
        info!("Testing p2p of {}:{:?}", type_name::<M>(), m);
        let to = *self
            .clients_list_minus_host::<M>()
            .iter()
            .next()
            .expect("No clients available to send message to");

        self.host.send_to(to, m.clone());

        self.assert_received_targeted_message(m, to);
    }

    /// Sends and receives and checks tha the message was ok.
    ///
    /// `client_sender`: which client to use as sender, or None to use the host.
    ///
    /// # Panics
    ///
    /// Fails if the assertions fail.
    pub fn broadcast_receive<M>(&mut self, msg: &M, client_sender: Option<usize>)
    where
        M: Message + Clone + Debug + PartialEq,
    {
        let sender = if let Some(i) = client_sender {
            assert!(i < self.clients.len(), "Invalid index for client.");
            &mut self.clients[i]
        } else {
            &mut self.host
        };

        let sender_name = sender.name.clone();
        sender.broadcast(msg.clone());
        info!("Sent message in broadcasting test");
        if client_sender.is_some() {
            self.assert_host_received(msg, client_sender, &sender_name);
        }

        self.assert_all_received_message(msg, client_sender, &sender_name);
    }

    fn assert_disconnected<M: Message>(&mut self) -> Result<(), String> {
        if is_connected::<M>(&self.host.app) {
            return Err("Host is still connected".to_string());
        }
        for client in &self.clients {
            if is_connected::<M>(&client.app) {
                return Err("Client is still connected".to_string());
            }
        }
        Ok(())
    }

    fn assert_host_uuid_reported<M: Message>(&mut self) {
        let host = self.host.uuid::<M>();
        assert_eq!(host, self.host.host_uuid::<M>());
        for c in &mut self.clients {
            assert_eq!(host, c.host_uuid::<M>());
        }
    }

    fn assert_host_is_in_destinations<M: Message>(&mut self) {
        let host = self.host.uuid::<M>();
        let tos = self.host.destinations::<M>();
        assert!(
            tos.contains(&host),
            "The host uuid {host} was not in the destination list",
        );
    }

    fn assert_clients_are_in_destinations<M: Message>(&mut self) {
        for c in &mut self.clients {
            let uuid = c.uuid::<M>();
            let tos = c.destinations::<M>();
            assert!(
                tos.contains(&uuid),
                "The client uuid {uuid} was not in the destination list",
            );
        }
    }

    fn assert_uuids_all_unique<M: Message>(&mut self) {
        let mut uuids = vec![];
        uuids.push(self.host.uuid::<M>());
        for c in &mut self.clients {
            let uuid = c.uuid::<M>();
            assert!(
                !uuids.contains(&uuid),
                "The client uuid {uuid} was not uniquely assigned.",
            );
            uuids.push(uuid);
        }
    }

    fn assert_known_host<M: Message>(&mut self) {
        let host = self.host.uuid::<M>();
        for c in &mut self.clients {
            assert_eq!(host, c.host_uuid::<M>());
        }
    }

    fn print_all_destinations<M: Message>(&mut self, prefix: &str) {
        info!("Host is {}", self.host.uuid::<M>());

        let mut destinations_collection: HashMap<ClientId, HashSet<ClientId>> = HashMap::new();
        destinations_collection.insert(self.host.uuid::<M>(), self.host.destinations::<M>());
        for c in &mut self.clients {
            destinations_collection.insert(c.uuid::<M>(), c.destinations::<M>());
        }

        for (collection_uuid, destinations) in &destinations_collection {
            let dests = destinations
                .iter()
                .map(|c| format!("{c}"))
                .collect::<Vec<_>>()
                .join(", ");
            info!(
                "{prefix} {} Destinations\n{collection_uuid} -> {dests}",
                type_name::<M>(),
            );
        }
    }

    fn assert_all_destinations<M: Message>(&mut self) {
        self.print_all_destinations::<M>("Asserting destinations after migration,");

        retry(|| {
            self.update_all();

            let mut destinations_collection: HashMap<ClientId, HashSet<ClientId>> = HashMap::new();
            destinations_collection.insert(self.host.uuid::<M>(), self.host.destinations::<M>());
            for c in &mut self.clients {
                destinations_collection.insert(c.uuid::<M>(), c.destinations::<M>());
            }

            for (collection_uuid, destinations) in destinations_collection {
                let uuid = self.host.uuid::<M>();
                if !destinations.contains(&uuid) {
                    return Err(format!(
                        "Host {uuid} not found in destinations of {collection_uuid}",
                    ));
                }
                for c in &mut self.clients {
                    let uuid = c.uuid::<M>();
                    if !destinations.contains(&uuid) {
                        return Err(format!(
                            "Client {uuid} not found in destinations of {collection_uuid}",
                        ));
                    }
                }
            }
            Ok(())
        })
        .unwrap();
    }

    fn assert_host_received<M>(
        &mut self,
        msg: &M,
        client_sender: Option<usize>,
        sender_name: &String,
    ) where
        M: Message + Clone + Debug + PartialEq,
    {
        retry(|| {
            let m = self.host.recv::<M>();
            if m.is_none() {
                return Err(format!(
                    "No message for {}, iteration {:?}, sender was {}",
                    self.host.name, client_sender, sender_name,
                ));
            }
            let m = m.ok_or("No message found".to_string())?;
            if *m != msg.clone() {
                return Err(format!(
                    "Wrong message for {}, iteration {:?}, sender was {}",
                    self.host.name, client_sender, sender_name,
                ));
            }
            Ok(())
        })
        .unwrap();
    }
    fn assert_all_received_message<M>(
        &mut self,
        msg: &M,
        client_sender: Option<usize>,
        sender_name: &str,
    ) where
        M: Message + Clone + Debug + PartialEq,
    {
        for (idx, c) in self.clients.iter_mut().enumerate() {
            let r = if let Some(i) = client_sender
                && i == idx
            {
                continue;
            } else {
                c
            };
            retry(|| {
                let m = r.recv::<M>();
                if m.is_none() {
                    return Err(format!(
                        "No message for {}, iteration {:?}, sender was {}",
                        r.name, client_sender, sender_name,
                    ));
                }
                let m = m.ok_or("No message found".to_string())?;
                if *m != msg.clone() {
                    return Err(format!(
                        "Wrong message for {}, iteration {:?}, sender was {}",
                        r.name, client_sender, sender_name,
                    ));
                }
                Ok(())
            })
            .unwrap();
        }
    }

    fn assert_received_targeted_message<M>(&mut self, m: &M, to: ClientId)
    where
        M: Message + Clone + Debug + PartialEq,
    {
        // Check that this client received this message.
        retry(|| {
            let rec_client = self.client_of_uuid::<M>(to);
            let rec_m = rec_client
                .recv::<M>()
                .ok_or(format!("Expected message on client uuid {to}"))?;
            if *m != *rec_m {
                return Err("Wrong message received".to_string());
            }
            Ok(())
        })
        .unwrap();

        // Check that all the other clients have NOT received this message
        for c in &mut self.clients {
            if c.uuid::<M>() == to {
                continue;
            }
            assert!(
                c.recv::<M>().is_none(),
                "A message was received on client uudi {to} but it should not have.",
            );
        }
    }

    fn clients_list_minus_host<M: Message>(&mut self) -> HashSet<ClientId> {
        let host = self.host.uuid::<M>();
        let mut tos = self.host.destinations::<M>();
        tos.remove(&host);
        tos
    }

    fn client_of_uuid<M: Message>(&mut self, uuid: ClientId) -> &mut TestApp {
        for c in &mut self.clients {
            let uuid_client = c.uuid::<M>();
            if uuid_client == uuid {
                return c;
            }
        }
        panic!("Client with uuid {uuid} not found");
    }

    pub fn update_all(&mut self) {
        self.host.app.update();
        for c in &mut self.clients {
            c.app.update();
        }
    }
}

pub struct TestApp {
    name: String,
    app: App,
}

impl TestApp {
    fn new(name: &str) -> Self {
        let mut app = App::new();
        app.add_plugins(SessionPlugin::<Msg>::default());
        app.add_plugins(SessionPlugin::<Msg2>::default());
        Self {
            name: name.to_string(),
            app,
        }
    }

    #[allow(clippy::needless_pass_by_value)]
    fn new_host(port: u16, compress: bool, key: Option<Vec<u8>>) -> Self {
        let mut test_app = Self::new("host");
        test_app
            .app
            .world_mut()
            .commands()
            .queue(SessionConnectCommand::<Msg>::from_config(
                SessionConfig::Direct {
                    addr: Some("127.0.0.1".parse().unwrap()),
                    port,
                    host: true,
                    compress,
                    key: key.clone(),
                    options: SessionOptions::default(),
                },
            ));
        test_app
            .app
            .world_mut()
            .commands()
            .queue(SessionConnectCommand::<Msg2>::from_config(
                SessionConfig::Direct {
                    addr: Some("127.0.0.1".parse().unwrap()),
                    port: port + 1,
                    host: true,
                    compress,
                    key: key.clone(),
                    options: SessionOptions::default(),
                },
            ));
        test_app.app.update();
        test_app.assert_connected_event();
        test_app.assert_channel_present();
        test_app
    }

    #[allow(clippy::needless_pass_by_value)]
    fn new_client(port: u16, compress: bool, key: Option<Vec<u8>>) -> Self {
        let mut test_app = Self::new("client");
        test_app
            .app
            .world_mut()
            .commands()
            .queue(SessionConnectCommand::<Msg>::from_config(
                SessionConfig::Direct {
                    addr: Some("127.0.0.1".parse().unwrap()),
                    port,
                    host: false,
                    compress,
                    key: key.clone(),
                    options: SessionOptions::default(),
                },
            ));
        test_app
            .app
            .world_mut()
            .commands()
            .queue(SessionConnectCommand::<Msg2>::from_config(
                SessionConfig::Direct {
                    addr: Some("127.0.0.1".parse().unwrap()),
                    port: port + 1,
                    host: false,
                    compress,
                    key: key.clone(),
                    options: SessionOptions::default(),
                },
            ));
        test_app.app.update();
        test_app.assert_connected_event();
        test_app.assert_channel_present();
        // Need to update the app again so the event is processed by the system.
        test_app.app.update();
        test_app
    }

    #[must_use]
    pub fn new_triple(port: u16) -> TestGroup {
        let host = Self::new_host(port, false, None);
        let mut client1 = Self::new_client(port, false, None);
        client1.name = "client1".to_string();
        let mut client2 = Self::new_client(port, false, None);
        client2.name = "client2".to_string();
        TestGroup {
            port,
            host,
            clients: vec![client1, client2],
        }
    }

    #[must_use]
    pub fn new_triple_comp_enc(port: u16) -> TestGroup {
        let host = Self::new_host(port, true, Some(KEY.to_vec()));
        let mut client1 = Self::new_client(port, true, Some(KEY.to_vec()));
        client1.name = "client1".to_string();
        let mut client2 = Self::new_client(port, true, Some(KEY.to_vec()));
        client2.name = "client2".to_string();
        TestGroup {
            port,
            host,
            clients: vec![client1, client2],
        }
    }

    pub fn is_host<M: Message>(&self) -> bool {
        self.app.world().resource::<Channel<M>>().is_host()
    }

    pub fn promote_new_host<M: Message>(&mut self, new_host: ClientId, port: u16) {
        let promote = SessionPromoteToHostCommand::<M>::new(new_host, Some(port));
        self.app.world_mut().commands().queue(promote);
        self.app.update();
    }

    pub fn broadcast<M: Message>(&mut self, m: M) {
        self.channel::<M>().broadcast(m);
    }

    pub fn send_to<M: Message>(&mut self, to: ClientId, m: M) {
        self.channel::<M>().send_to(to, m);
    }

    /// # Panics
    ///
    /// Panics if channel does not exist.
    pub fn channel<M: Message>(&mut self) -> Mut<'_, Channel<M>> {
        let Some(c) = self.app.world_mut().get_resource_mut::<Channel<M>>() else {
            panic!(
                "Resource Channel<{}> not found on {}",
                type_name::<M>(),
                self.name
            );
        };
        c
    }

    pub fn uuid<M: Message>(&self) -> ClientId {
        self.get_channel::<M>().uuid()
    }

    /// # Panics
    ///
    /// Panics if host does not exist.
    pub fn host_uuid<M: Message>(&mut self) -> ClientId {
        self.get_channel::<M>().host_uuid().unwrap()
    }

    pub fn destinations<M: Message>(&mut self) -> HashSet<ClientId> {
        self.get_channel::<M>().destinations()
    }

    pub fn recv<M: Message + Clone>(&mut self) -> Option<Box<M>> {
        self.app.update();
        let events = self
            .app
            .world_mut()
            .resource_mut::<Messages<MessageReceivedEvent<M>>>();
        let mut cursor = events.get_cursor();
        let e = cursor.read(&events).next();
        e.map(|e| e.message.clone())
    }

    fn disconnect_all(&mut self) {
        self.app
            .world_mut()
            .commands()
            .queue(SessionDisconnectCommand::<Msg>::default());
        self.app
            .world_mut()
            .commands()
            .queue(SessionDisconnectCommand::<Msg2>::default());
        self.app.update();
        self.assert_disconnected_event();
    }

    fn disconnect_one<M: Message>(&mut self) {
        self.app
            .world_mut()
            .commands()
            .queue(SessionDisconnectCommand::<M>::default());
        self.app.update();
        self.assert_disconnected_event_one::<M>();
    }

    fn assert_disconnected_event(&mut self) {
        self.assert_event::<SessionDisconnectedEvent<Msg>>("disconnection check both");
        self.assert_event::<SessionDisconnectedEvent<Msg2>>("disconnection check both");
    }

    fn assert_disconnected_event_one<M: Message>(&mut self) {
        self.assert_event::<SessionDisconnectedEvent<M>>("disconnection check single");
    }

    fn assert_connected_event(&mut self) {
        self.assert_event::<SessionConnectedEvent<Msg>>("connect check both");
        self.assert_event::<SessionConnectedEvent<Msg2>>("connect check both");
    }

    fn assert_event<E: BevyMessage>(&mut self, msg: &str) {
        // let max_retries = 10;
        // let mut retries = 0;
        // while retries < max_retries && self.has_event::<E>() {
        //     retries += 1;
        //     self.app.update();
        // }
        // assert!(
        //     retries != max_retries,
        //     "Expected event {:?} on {}",
        //     type_name::<E>(),
        //     self.name.as_str()
        // );
        assert!(
            self.has_event::<E>(),
            "Expected event {:?} on {}, {msg}",
            type_name::<E>(),
            self.name.as_str()
        );
    }

    fn has_event<E: BevyMessage>(&mut self) -> bool {
        let events = self.app.world_mut().resource_mut::<Messages<E>>();
        let mut cursor = events.get_cursor();
        cursor.read(&events).next().is_some()
    }

    fn get_channel<M: Message>(&self) -> &Channel<M> {
        retry(|| {
            self.app.world().get_resource::<Channel<M>>().ok_or(format!(
                "Expected Channel<{}> on {}",
                type_name::<M>(),
                self.name.as_str()
            ))
        })
        .unwrap()
    }

    fn assert_channel_present(&mut self) {
        self.get_channel::<Msg>();
        self.get_channel::<Msg2>();
    }
}

impl Drop for TestApp {
    fn drop(&mut self) {
        self.disconnect_all();
    }
}

fn is_connected<M: Message>(app: &App) -> bool {
    app.world().get_resource::<Channel<M>>().is_some()
}

fn retry<T, E, F: FnMut() -> Result<T, E>>(mut f: F) -> Result<T, E> {
    let mut ct = 0;
    loop {
        let res = f();
        match res {
            Ok(t) => return Ok(t),
            Err(e) => {
                if ct >= 50 {
                    return Err(e);
                }
            }
        }
        sleep(Duration::from_millis(100));
        ct += 1;
    }
}