simple-pub-sub 0.1.1

use crate::message;
use crate::stream;
use log::{error, info, trace};
use std::io::ErrorKind;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream,
};

type Callback = fn(String, Vec<u8>);
#[derive(Debug)]
pub struct Client {
    pub server: String,
    pub port: u16,
    stream: Option<TcpStream>,
    callback: Option<Callback>,
}

/// default implementation for callback function
pub fn on_message(topic: String, message: Vec<u8>) {
    match String::from_utf8(message.clone()) {
        Ok(msg_str) => {
            info!("topic: {} message: {}", topic, msg_str);
        }
        Err(_) => {
            info!("topic: {} message: {:?}", topic, message);
        }
    };
}

impl Client {
    pub fn new(server: String, port: u16) -> Client {
        Client {
            server,
            port,
            stream: None,
            callback: None,
        }
    }

    pub fn on_message(&mut self, callback: Callback) {
        self.callback = Some(callback)
    }

    pub async fn connect(&mut self) -> Result<(), tokio::io::Error> {
        let server_url: String = format!("{}:{}", self.server, self.port);
        let stream = TcpStream::connect(server_url).await?;
        self.stream = Some(stream);
        Ok(())
    }

    pub async fn post(self, msg: message::Msg) -> Result<Vec<u8>, tokio::io::Error> {
        match self.stream {
            Some(mut s) => {
                match s.write_all(&msg.bytes()).await {
                    Ok(_) => {}
                    Err(e) => {
                        error!("could not send the data to the server: {}", e.to_string());
                        return Err(e);
                    }
                }
                let mut buf: Vec<u8>;
                buf = vec![0; 8];
                trace!("reading the ack:");
                let _ = match s.read(&mut buf).await {
                    Ok(n) => n,
                    Err(e) => {
                        error!("could not retrive ack for published packet");
                        return Err(e);
                    }
                };
                return Ok(buf);
            }
            None => {
                return Err(tokio::io::Error::new(
                    tokio::io::ErrorKind::Other,
                    "client not connected yet.",
                ));
            }
        };
    }

    pub async fn publish(self, topic: String, message: Vec<u8>) -> Result<(), tokio::io::Error> {
        let msg: message::Msg = message::Msg::new(message::PktType::PUBLISH, topic, Some(message));
        trace!("msg: {:?}", msg);

        let buf = match self.post(msg).await {
            Ok(resp) => resp,
            Err(e) => return Err(e),
        };

        trace!("the raw buffer is: {:?}", buf);
        let resp_: message::Header = match message::Header::from_vec(buf) {
            Ok(resp) => resp,
            Err(e) => {
                return Err(tokio::io::Error::new(ErrorKind::Other, format!("{:?}", e)));
            }
        };
        trace!("{:?}", resp_);

        Ok(())
    }

    pub async fn query(self, topic: String) -> Result<String, tokio::io::Error> {
        let msg: message::Msg = message::Msg::new(
            message::PktType::QUERY,
            topic,
            Some("a".to_string().as_bytes().to_vec()),
        );
        trace!("msg: {:?}", msg);

        match self.stream {
            Some(mut s) => match s.write_all(&msg.bytes()).await {
                Ok(_) => {
                    match stream::read_message(&mut s).await {
                        Ok(msg) => {
                            match String::from_utf8(msg.message.clone()) {
                                Ok(msg_str) => {
                                    return Ok(msg_str);
                                }
                                Err(e) => {
                                    return Err(tokio::io::Error::new(
                                        tokio::io::ErrorKind::Other,
                                        e.to_string(),
                                    ));
                                }
                            };
                        }
                        Err(e) => {
                            return Err(e);
                        }
                    };
                }
                Err(e) => {
                    return Err(tokio::io::Error::new(
                        tokio::io::ErrorKind::Other,
                        e.to_string(),
                    ));
                }
            },
            None => {
                return Err(tokio::io::Error::new(
                    tokio::io::ErrorKind::Other,
                    "client not connected yet.".to_string(),
                ));
            }
        };
    }

    pub async fn subscribe(self, topic: String) -> Result<(), tokio::io::Error> {
        match self.stream {
            Some(mut s) => {
                let msg: message::Msg = message::Msg::new(message::PktType::SUBSCRIBE, topic, None);
                trace!("msg: {:?}", msg);

                match s.write_all(&msg.bytes()).await {
                    Ok(_) => {}
                    Err(e) => {
                        error!("could not send the data to the server: {}", e.to_string());
                        return Err(tokio::io::Error::new(
                            tokio::io::ErrorKind::Other,
                            e.to_string(),
                        ));
                    }
                };
                loop {
                    match stream::read_message(&mut s).await {
                        Ok(m) => match self.callback {
                            Some(fn_) => {
                                fn_(m.topic, m.message);
                            }
                            None => on_message(m.topic, m.message),
                        },
                        Err(e) => {
                            error!("could not read message: {}", e.to_string());
                            break;
                        }
                    };
                }

                Ok(())
            }
            None => {
                return Err(tokio::io::Error::new(
                    tokio::io::ErrorKind::Other,
                    "client not connected yet.".to_string(),
                ));
            }
        }
    }
}