tubes 0.6.5

use std::collections::HashMap;
use std::{collections::HashSet, fmt::Debug, string::FromUtf8Error, thread::sleep, time::Duration};
use tubes::ClientId;
use tubes::prelude::*;

#[derive(Clone, Debug, PartialEq)]
pub struct Msg(pub String);

impl From<String> for Msg {
    fn from(value: String) -> Self {
        Self(value)
    }
}

impl TryFrom<&[u8]> for Msg {
    type Error = FromUtf8Error;

    fn try_from(value: &[u8]) -> std::result::Result<Self, Self::Error> {
        Ok(Msg(String::from_utf8(value.to_vec())?))
    }
}

impl TryFrom<Msg> for Vec<u8> {
    type Error = ();

    fn try_from(value: Msg) -> std::result::Result<Self, Self::Error> {
        Ok(value.0.into())
    }
}

pub struct TestSession {
    server: Session,
    nodes: Vec<Session>,
    promoting: Option<ClientId>,
}

impl TestSession {
    #[must_use]
    pub fn new_pair(port: u16) -> Self {
        let mut res = Self {
            server: Self::new_server(port),
            nodes: vec![Self::new_client(port)],
            promoting: None,
        };
        res.connect();
        res
    }

    #[must_use]
    pub fn new_triple(port: u16) -> Self {
        let mut res = Self {
            server: Self::new_server(port),
            nodes: vec![Self::new_client(port), Self::new_client(port)],
            promoting: None,
        };
        res.connect();
        res
    }

    #[must_use]
    pub fn new_quad(port: u16) -> Self {
        let mut res = Self {
            server: Self::new_server(port),
            nodes: vec![
                Self::new_client(port),
                Self::new_client(port),
                Self::new_client(port),
            ],
            promoting: None,
        };
        res.connect();
        res
    }

    #[cfg(all(feature = "encryption", feature = "compression"))]
    #[must_use]
    pub fn new_pair_pack(port: u16) -> Self {
        let mut res = Self {
            server: Self::new_server_pack(port),
            nodes: vec![Self::new_client_pack(port)],
            promoting: None,
        };
        res.connect();
        res
    }

    #[cfg(all(feature = "encryption", feature = "compression"))]
    #[must_use]
    pub fn new_triple_pack(port: u16) -> Self {
        let mut res = Self {
            server: Self::new_server_pack(port),
            nodes: vec![Self::new_client_pack(port), Self::new_client_pack(port)],
            promoting: None,
        };
        res.connect();
        res
    }

    #[cfg(all(feature = "encryption", feature = "compression"))]
    #[must_use]
    pub fn new_quad_pack(port: u16) -> Self {
        let mut res = Self {
            server: Self::new_server_pack(port),
            nodes: vec![
                Self::new_client_pack(port),
                Self::new_client_pack(port),
                Self::new_client_pack(port),
            ],
            promoting: None,
        };
        res.connect();
        res
    }

    pub fn stop(&mut self) {
        for c in &mut self.nodes {
            c.stop();
        }
        self.server.stop();
    }

    #[must_use]
    pub fn server_uuid(&self) -> ClientId {
        self.server.uuid()
    }

    #[must_use]
    pub fn clients(&self) -> HashSet<ClientId> {
        self.server.clients()
    }

    pub fn client_leave(&mut self, uuid: ClientId) {
        // Dropping the client will also close it.
        self.nodes.retain(|c| c.uuid() != uuid);
    }

    /// Clean out any receiving queue from all the nodes, allowing for a clear
    /// test trigger of a specific message without pending scenarios.
    ///
    /// # Panics
    ///
    /// Panics if message pipe is broken.
    pub fn exhaust(&mut self) {
        // TODO: this is time sensitive
        sleep(Duration::from_millis(100));
        // During promotion always exhaust the promoting client first
        if let Some(p) = self.promoting {
            for c in &mut self.nodes {
                if p == c.uuid() {
                    while c.read().unwrap().is_some() {}
                    self.promoting = None;
                }
            }
        }
        while self.server.read().unwrap().is_some() {}
        for c in &mut self.nodes {
            while c.read().unwrap().is_some() {}
        }
    }

    /// # Panics
    ///
    /// Panics if message is not convertible.
    pub fn server_broadcast(&mut self, m: Msg) {
        self.server.broadcast(m.try_into().unwrap()).unwrap();
    }

    /// # Panics
    ///
    /// Panics if message is not convertible.
    pub fn server_send_to(&mut self, uuid: ClientId, m: Msg) {
        self.server.send_to(uuid, m.try_into().unwrap()).unwrap();
    }

    /// # Panics
    ///
    /// Panics if message is not convertible.
    pub fn first_client_send_to_server(&mut self, m: Msg) -> Option<ClientId> {
        let c = self.nodes.iter_mut().next()?;
        c.send_to(self.server.uuid(), m.try_into().unwrap())
            .unwrap();
        Some(c.uuid())
    }

    /// # Panics
    ///
    /// Panics if message is not convertible.
    pub fn first_client_broadcast(&mut self, m: Msg) -> Option<ClientId> {
        let c = self.nodes.iter_mut().next()?;
        c.broadcast(m.try_into().unwrap()).unwrap();
        Some(c.uuid())
    }

    /// # Panics
    ///
    /// Panics if clients are not found.
    pub fn promote_to_host(&mut self, client: ClientId, port: u16) {
        println!("TEST /// Before migration");
        self.print_all_destinations();
        self.assert_all_destinations();

        self.promoting = Some(client);
        self.server.promote_to_host(client, Some(port));
        // By using the session first, the reconnection is being triggered.
        // The reconnection does not trigger until at least one interaction from
        // the user happens.
        self.exhaust();
        self.assert_is_server(client);

        // TODO: time sensitive
        std::thread::sleep(Duration::from_millis(100));

        println!("TEST /// After migration");
        self.print_all_destinations();

        // Swap the promoted client with the server position or the testing
        // will not work afterwards.
        let promoted_client = self
            .nodes
            .iter_mut()
            .find(|c| c.uuid() == client)
            .ok_or(format!("Client not found after mitration: {client}"))
            .unwrap();
        std::mem::swap(&mut self.server, promoted_client);

        self.assert_all_destinations();
    }

    /// # Panics
    ///
    /// Panics if message is not convertible.
    pub fn assert_server_received_broacast(&mut self, from: ClientId, m: Msg) {
        let m = MessageData::Broadcast {
            from,
            data: m.try_into().unwrap(),
        };
        retry(move || Self::assert_received_message(&mut self.server, &m)).unwrap();
    }

    /// # Panics
    ///
    /// Panics if message is not convertible.
    pub fn assert_all_clients_received_broacast(&mut self, from: ClientId, m: Msg) {
        let m = MessageData::Broadcast {
            from,
            data: m.try_into().unwrap(),
        };
        for c in &mut self.nodes {
            retry(|| Self::assert_received_message(c, &m)).unwrap();
        }
    }

    /// # Panics
    ///
    /// Panics if message is not convertible.
    pub fn assert_only_client_received(&mut self, from: ClientId, uuid: ClientId, m: Msg) {
        let m = MessageData::Send {
            from,
            to: uuid,
            data: m.try_into().unwrap(),
        };
        for c in &mut self.nodes {
            if c.uuid() == uuid {
                retry(|| Self::assert_received_message(c, &m)).unwrap();
            }
        }
        // Iterate again: checking when a message does not exist has to be done
        // in a single shot, so make sure to do it only after the positive case
        // above has been already checked and wait a little for the other
        // clients to catch up.
        for c in &mut self.nodes {
            if c.uuid() != uuid {
                assert!(c.read().unwrap().is_none());
            }
        }
    }

    /// # Panics
    ///
    /// Panics if message is not convertible.
    pub fn assert_only_server_received(&mut self, from: ClientId, m: Msg) {
        let m = MessageData::Send {
            from,
            to: self.server.uuid(),
            data: m.try_into().unwrap(),
        };
        retry(|| Self::assert_received_message(&mut self.server, &m)).unwrap();
        // Iterate again: checking when a message does not exist has to be done
        // in a single shot, so make sure to do it only after the positive case
        // above has been already checked and wait a little for the other
        // clients to catch up.
        for c in &mut self.nodes {
            assert!(c.read().unwrap().is_none());
        }
    }

    /// # Panics
    ///
    /// Panics if no messages are there.
    pub fn assert_client_left(&mut self, uuid: ClientId) {
        let m = MessageData::ClientLeft(uuid);
        retry(|| Self::assert_received_message(&mut self.server, &m)).unwrap();
        for c in &mut self.nodes {
            retry(|| Self::assert_received_message(c, &m)).unwrap();
        }
        retry(|| {
            if self.server.clients().contains(&uuid) {
                return Err(format!("Client has not left {uuid}"));
            }
            Ok(())
        })
        .unwrap();
    }

    /// # Panics
    ///
    /// Panics if not connected.
    pub fn assert_is_server(&self, uuid: ClientId) {
        retry(|| {
            let client_is_server = self
                .nodes
                .iter()
                .find(|c| c.uuid() == uuid)
                .is_some_and(Session::is_server);
            let server_is_server = uuid == self.server.uuid();
            if !server_is_server && !client_is_server {
                return Err(format!("{uuid} is not server"));
            }
            Ok(())
        })
        .unwrap();
    }

    fn print_all_destinations(&mut self) {
        println!("Host is {}", self.server.uuid());

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

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

    fn assert_all_destinations(&mut self) {
        retry(|| {
            let mut destinations_collection: HashMap<ClientId, HashSet<ClientId>> = HashMap::new();
            destinations_collection.insert(self.server.uuid(), self.server.clients());
            for c in &mut self.nodes {
                destinations_collection.insert(c.uuid(), c.clients());
            }

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

    fn connect(&mut self) {
        self.server.start().unwrap();
        assert!(self.server.is_connected());
        for c in &mut self.nodes {
            c.start().unwrap();
            assert!(c.is_connected());
        }

        retry(|| {
            let len1 = self.server.clients().len();
            let len2 = self.nodes.len() + 1;
            if len1 != len2 {
                return Err(format!("Client lengths not matching: {len1} != {len2}"));
            }
            for uuid in self.nodes.iter().map(Session::uuid) {
                if !self.server.clients().contains(&uuid) {
                    return Err(format!("Client not in server list: {uuid}"));
                }
            }
            Ok(())
        })
        .unwrap();

        let server_uuid = self.server_uuid();
        for c in &mut self.nodes {
            retry(|| {
                let Some(suuid) = c.server_uuid() else {
                    return Err(format!(
                        "Client {} did not receive the server uuid",
                        c.uuid()
                    ));
                };
                if server_uuid != suuid {
                    return Err(format!(
                        "Client {} did receive the wrong server uuid {} != {}",
                        c.uuid(),
                        server_uuid,
                        suuid
                    ));
                }
                Ok(())
            })
            .unwrap();
        }

        self.assert_all_destinations();
    }

    fn new_server(port: u16) -> Session {
        let s = Session::new_server(format!("127.0.0.1:{port}").as_str().into());
        assert!(s.is_server());
        assert!(!s.is_connected());
        s
    }

    fn new_client(port: u16) -> Session {
        let s = Session::new_client(format!("127.0.0.1:{port}").as_str().into());
        assert!(!s.is_server());
        assert!(!s.is_connected());
        s
    }

    #[cfg(all(feature = "encryption", feature = "compression"))]
    fn new_server_pack(port: u16) -> Session {
        let config = Config {
            address: "127.0.0.1".parse().ok(),
            port,
            versions: Versions::default(),
            accept_timeout: Duration::default(),
            compress: true,
            key: Some(Self::key()),
        };
        Session::new_server(config)
    }

    #[cfg(all(feature = "encryption", feature = "compression"))]
    fn new_client_pack(port: u16) -> Session {
        let config = Config {
            address: "127.0.0.1".parse().ok(),
            port,
            versions: Versions::default(),
            accept_timeout: Duration::default(),
            compress: true,
            key: Some(Self::key()),
        };
        Session::new_client(config)
    }

    #[cfg(feature = "encryption")]
    fn key() -> Vec<u8> {
        vec![0; 32]
    }

    fn assert_received_message(s: &mut Session, m: &MessageData) -> Result<(), String> {
        let m = Self::debug_message(m);
        let msg = s.read().unwrap().ok_or(format!(
            "No message available on {}, expected {}",
            s.uuid(),
            m
        ))?;
        let msg = Self::debug_message(&msg);
        if msg != m {
            return Err(format!(
                "Wrong message, expected: \n{m}\n returned: \n{msg}",
            ));
        }
        Ok(())
    }

    fn debug_message(m: &MessageData) -> String {
        match m {
            MessageData::Broadcast { from, data: m } => format!("Broadcast({from}, {m:?})"),
            MessageData::Send {
                from,
                to: uuid,
                data: m,
            } => format!("Send({from}, {uuid}, {m:?})"),
            MessageData::ClientJoined(uuid) => format!("ClientJoined({uuid})"),
            MessageData::ClientLeft(uuid) => format!("ClientLeft({uuid})"),
        }
    }
}

impl Drop for TestSession {
    fn drop(&mut self) {
        self.stop();
    }
}

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