1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use futures::FutureExt;
use std::{
    future::Future,
    io::{BufReader, Read},
    process::{Command, Stdio},
};

pub use norgopolis_protos::client_communication::MessagePack;
use norgopolis_protos::client_communication::{forwarder_client::ForwarderClient, Invocation};

use serde::de::DeserializeOwned;
use tonic::{transport::Channel, Request, Response, Status, Streaming};

pub struct ConnectionHandle(ForwarderClient<Channel>);

impl ConnectionHandle {
    pub fn invoke_raw(
        &mut self,
        module: String,
        function_name: String,
        args: Option<MessagePack>,
    ) -> impl Future<Output = Result<Response<Streaming<MessagePack>>, Status>> + '_ {
        self.0.forward(Request::new(Invocation {
            module,
            function_name,
            args,
        }))
    }

    pub async fn invoke_raw_callback<F>(
        &mut self,
        module: String,
        function_name: String,
        args: Option<MessagePack>,
        callback: F,
    ) -> anyhow::Result<()>
    where
        F: Fn(MessagePack),
    {
        self.invoke_raw(module, function_name, args)
            .then(|response| async move {
                let mut response = response?.into_inner();

                while let Some(data) = response.message().await? {
                    callback(data);
                }

                Ok::<(), anyhow::Error>(())
            })
            .await?;

        Ok(())
    }

    pub async fn invoke<TargetStruct, F>(
        &mut self,
        module: String,
        function_name: String,
        args: Option<MessagePack>,
        callback: F,
    ) -> anyhow::Result<()>
    where
        F: Fn(TargetStruct),
        TargetStruct: DeserializeOwned,
    {
        self.invoke_raw(module, function_name, args)
            .then(|response| async move {
                let mut response = response?.into_inner();

                while let Some(data) = response.message().await? {
                    callback(rmp_serde::from_slice::<TargetStruct>(data.data.as_slice()).unwrap());
                }

                Ok::<(), anyhow::Error>(())
            })
            .await?;

        Ok(())
    }

    pub async fn invoke_collect<TargetStruct>(
        &mut self,
        module: String,
        function_name: String,
        args: Option<MessagePack>,
    ) -> anyhow::Result<Vec<TargetStruct>>
    where
        TargetStruct: DeserializeOwned,
    {
        let response = self.invoke_raw(module, function_name, args).await;

        let mut response = response?.into_inner();
        let mut result: Vec<TargetStruct> = Vec::new();

        while let Some(data) = response.message().await? {
            result.push(rmp_serde::from_slice::<TargetStruct>(data.data.as_slice()).unwrap());
        }

        Ok(result)
    }
}

pub async fn connect(ip: &String, port: &String) -> anyhow::Result<ConnectionHandle> {
    Ok(ConnectionHandle(
        match ForwarderClient::connect("http://".to_string() + ip + ":" + port).await {
            Ok(connection) => connection,
            Err(err) => {
                if cfg!(feature = "autostart-server") {
                    let mut command = Command::new("norgopolis-server")
                        .stdout(Stdio::piped())
                        .spawn()?;

                    if let Some(stdout) = command.stdout.take() {
                        let mut buffer = BufReader::new(stdout);
                        let mut str = [b' '; 5];

                        buffer.read_exact(&mut str)?;

                        if str.starts_with(b"ready") {
                            return Ok(ConnectionHandle(
                                ForwarderClient::connect("http://".to_string() + ip + ":" + port)
                                    .await?,
                            ));
                        }
                    }
                }

                return Err(err.into());
            }
        },
    ))
}

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

    #[tokio::test]
    async fn establish_connection() {
        connect(&"127.0.0.1".into(), &"62020".into())
            .await
            .unwrap()
            .invoke(
                "test-module".to_string(),
                "func-name".to_string(),
                None,
                |response: (String,)| println!("{}", response.0),
            )
            .await
            .unwrap();
    }
}