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
153
154
155
156
157
158
159

use std::convert::TryInto;
use std::result::Result as StdResult;
use std::time::Duration;

use crate::broker::{ReconnectCondition, ReconnectConfig};
use crate::client::WebSocketClient;
use crate::error::{Error, Result};

use url::Url;

#[derive(Debug, Clone)]
struct WebSocketClientBuilderInner {
    url: Url,
    reconnect: ReconnectConfig,
}

/// Builder for [`WebSocketClient`].
#[derive(Debug, Clone)]
pub struct WebSocketClientBuilder {
    inner: Result<WebSocketClientBuilderInner>,
}

trait ResultExt<T, E> {
    fn and_then_mut<F>(&mut self, op: F)
    where
        F: FnOnce(&mut T) -> StdResult<(), E>;
}

impl<T, E> ResultExt<T, E> for StdResult<T, E> {
    fn and_then_mut<F>(&mut self, f: F)
    where
        F: FnOnce(&mut T) -> StdResult<(), E>,
    {
        if let Ok(x) = self {
            if let Err(e) = f(x) {
                *self = Err(e);
            }
        }
    }
}

impl WebSocketClientBuilder {
    /// Creates a new builder instance with `url`.
    ///
    /// All configurations are set to default.
    pub fn new<T>(url: T) -> Self
    where
        T: TryInto<Url>,
        T::Error: Into<Error>,
    {
        let inner = url
            .try_into()
            .map_err(Into::into)
            .map(|url| WebSocketClientBuilderInner {
                url,
                reconnect: ReconnectConfig::default(),
            });

        WebSocketClientBuilder { inner }
    }

    /// Creates a new builder instance with the given host name `host`.
    ///
    /// This method configures the builder with a URL of the form `wss://{host}/streaming`.
    /// All other configurations are set to default.
    pub fn with_host<S>(host: S) -> Self
    where
        S: AsRef<str>,
    {
        let url = format!("wss://{}/streaming", host.as_ref());
        WebSocketClientBuilder::new(url.as_str())
    }

    /// Sets an API token.
    ///
    /// This method appends the given token as the `i` query parameter to the URL.
    pub fn token<S>(&mut self, token: S) -> &mut Self
    where
        S: AsRef<str>,
    {
        self.query("i", token)
    }

    /// Specifies additional query parameters for the URL.
    pub fn query<S1, S2>(&mut self, key: S1, value: S2) -> &mut Self
    where
        S1: AsRef<str>,
        S2: AsRef<str>,
    {
        self.inner.and_then_mut(|inner| {
            inner
                .url
                .query_pairs_mut()
                .append_pair(key.as_ref(), value.as_ref());
            Ok(())
        });
        self
    }

    /// Sets whether or not to enable automatic reconnection.
    ///
    /// Automatic reconnection is enabled by default (as per [`Default`][default] implementation for
    /// [`ReconnectConfig`]), and you can disable it with `.auto_reconnect(false)`.
    ///
    /// [default]: std::default::Default
    pub fn auto_reconnect(&mut self, enable: bool) -> &mut Self {
        if enable {
            self.reconnect_condition(ReconnectCondition::unexpected_reset())
        } else {
            self.reconnect_condition(ReconnectCondition::never())
        }
    }

    /// Sets an interval duration of automatic reconnection in seconds.
    pub fn reconnect_secs(&mut self, secs: u64) -> &mut Self {
        self.inner.and_then_mut(|inner| {
            inner.reconnect.interval = Duration::from_secs(secs);
            Ok(())
        });
        self
    }

    /// Sets an interval duration of automatic reconnection.
    pub fn reconnect_interval(&mut self, interval: Duration) -> &mut Self {
        self.inner.and_then_mut(|inner| {
            inner.reconnect.interval = interval;
            Ok(())
        });
        self
    }

    /// Specifies the condition for reconnecting.
    pub fn reconnect_condition(&mut self, condition: ReconnectCondition) -> &mut Self {
        self.inner.and_then_mut(|inner| {
            inner.reconnect.condition = condition;
            Ok(())
        });
        self
    }

    /// Specifies whether to re-send messages that may have failed to be sent when reconnecting.
    pub fn reconnect_retry_send(&mut self, enable: bool) -> &mut Self {
        self.inner.and_then_mut(|inner| {
            inner.reconnect.retry_send = enable;
            Ok(())
        });
        self
    }

    /// Finish this builder instance and connect to Misskey using this configuration.
    pub async fn connect(&self) -> Result<WebSocketClient> {
        let WebSocketClientBuilderInner { url, reconnect } = match self.inner.clone() {
            Err(e) => return Err(e),
            Ok(inner) => inner,
        };

        WebSocketClient::connect_with_config(url, reconnect).await
    }
}