1use std::time::{Duration, Instant};
2
3use log::{debug, error, info, trace, warn};
4use serde::de::DeserializeOwned;
5
6use whiskers_core::{Age, Direction, EmbedError, Embedder, Guard, GuardError, Household, IconChooser, IconError, IconId, MemorySnapshot, RankError, Ranker, Verdict};
7
8use whiskers_core::wire::{CheckReply, CheckRequest, EmbedReply, EmbedRequest, IconReply, IconRequest, RankReply, RankRequest};

11pub struct RemoteGuard {
12    base_url: String,
13    agent: ureq::Agent,
14}

16impl RemoteGuard {
17    pub fn new(base_url: impl Into<String>) -> Self {
18        debug!("remote guard client, timeout 15 s");
19        let agent = ureq::Agent::config_builder()
20            .timeout_global(Some(Duration::from_secs(15)))
21            .http_status_as_error(false)
22            .build()
23            .into();
24        Self { base_url: base_url.into().trim_end_matches('/').to_owned(), agent }
25    }
26}
27
28impl Guard for RemoteGuard {
29    fn check(&self, direction: Direction, age: Age, text: &str) -> Result<Verdict, GuardError> {
30        debug!("remote guard check {direction:?} for age {}: {} chars", age.years(), text.len());
31        let started = Instant::now();
32        let body = serde_json::to_string(&CheckRequest { direction, age: Some(age), text: text.to_owned() }).map_err(|e| {
33            error!("guard request does not serialize: {e}");
34            GuardError(e.to_string())
35        })?;
36        let mut resp = self
37            .agent
38            .post(format!("{}/check", self.base_url))
39            .header("content-type", "application/json")
40            .send(body)
41            .map_err(|e| {
42                error!("guard service unreachable after {} ms: {e}", started.elapsed().as_millis());
43                GuardError(e.to_string())
44            })?;
45        let status = resp.status();
46        let text = resp.body_mut().read_to_string().map_err(|e| {
47            error!("guard reply unreadable (status {status}): {e}");
48            GuardError(e.to_string())
49        })?;
50        info!("guard service answered {status} in {} ms, {} bytes", started.elapsed().as_millis(), text.len());
51        if !status.is_success() {
52            error!("guard service refused with {status}");
53            return Err(GuardError(format!("guard service answered {status}")));
54        }
55        match serde_json::from_str::<CheckReply>(&text).map_err(|e| {
56            error!("guard reply does not parse: {e}");
57            GuardError(e.to_string())
58        })? {
59            CheckReply::Verdict(v) => Ok(v),
60            CheckReply::Unavailable(why) => {
61                warn!("guard service could not decide: {why}");
62                Err(GuardError(why))
63            }
64        }
65    }
66}

69fn post(base: &str, path: &str, body: String, seconds: u64) -> Result<String, String> {
70    let started = Instant::now();
71    debug!("POST {path}: {} byte body, timeout {seconds} s", body.len());
72    let agent: ureq::Agent = ureq::Agent::config_builder()
73        .timeout_global(Some(Duration::from_secs(seconds)))
74        .http_status_as_error(false)
75        .build()
76        .into();
77    let mut resp = agent
78        .post(format!("{}{path}", base.trim_end_matches('/')))
79        .header("content-type", "application/json")
80        .send(body)
81        .map_err(|e| {
82            warn!("POST {path} failed after {} ms: {e}", started.elapsed().as_millis());
83            e.to_string()
84        })?;
85    let status = resp.status();
86    let text = resp.body_mut().read_to_string().map_err(|e| {
87        warn!("POST {path} body unreadable (status {status}): {e}");
88        e.to_string()
89    })?;
90    if !status.is_success() {
91        warn!("POST {path} answered {status} in {} ms", started.elapsed().as_millis());
92        return Err(format!("service answered {status}"));
93    }
94    debug!("POST {path} answered {status} in {} ms, {} bytes", started.elapsed().as_millis(), text.len());
95    Ok(text)
96}

99fn parse<T: DeserializeOwned>(what: &str, text: &str) -> Result<T, String> {
100    serde_json::from_str(text).map_err(|e| {
101        warn!("{what} reply ({} bytes) does not parse: {e}", text.len());
102        e.to_string()
103    })
104}

108pub struct RemoteEmbedder {
109    base_url: String,
110}

112impl RemoteEmbedder {
113    pub fn new(base_url: impl Into<String>) -> Self {
114        Self { base_url: base_url.into() }
115    }
116}
117
118impl Embedder for RemoteEmbedder {
119    fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbedError> {
120        debug!("remote embed of {} text(s)", texts.len());
121        let body = serde_json::to_string(&EmbedRequest { texts: texts.to_vec() }).map_err(|e| EmbedError(e.to_string()))?;
122        let text = post(&self.base_url, "/embed", body, 6).map_err(EmbedError)?;
123        match parse::<EmbedReply>("embed", &text).map_err(EmbedError)? {
124            EmbedReply::Vectors(v) if v.len() == texts.len() => Ok(v),
125            EmbedReply::Vectors(v) => {
126                warn!("embed: asked for {} vectors, got {}", texts.len(), v.len());
127                Err(EmbedError(format!("asked for {} vectors, got {}", texts.len(), v.len())))
128            }
129            EmbedReply::Unavailable(why) => {
130                warn!("embed: service unavailable: {why}");
131                Err(EmbedError(why))
132            }
133        }
134    }
135}

138pub struct RemoteRanker {
139    base_url: String,
140}

142impl RemoteRanker {
143    pub fn new(base_url: impl Into<String>) -> Self {
144        Self { base_url: base_url.into() }
145    }
146}
147
148impl Ranker for RemoteRanker {
149    fn rank(&self, query: &str, candidates: &[String]) -> Result<Vec<f32>, RankError> {
150        debug!("remote rank of {} candidate(s)", candidates.len());
151        let body = serde_json::to_string(&RankRequest { query: query.to_owned(), candidates: candidates.to_vec() })
152            .map_err(|e| RankError(e.to_string()))?;
153        let text = post(&self.base_url, "/rerank", body, 8).map_err(RankError)?;
154        match parse::<RankReply>("rerank", &text).map_err(RankError)? {
155            RankReply::Probabilities(p) if p.len() == candidates.len() => Ok(p),
156            RankReply::Probabilities(p) => {
157                warn!("rerank: asked about {} candidates, got {}", candidates.len(), p.len());
158                Err(RankError(format!("asked about {} candidates, got {}", candidates.len(), p.len())))
159            }
160            RankReply::Unavailable(why) => {
161                warn!("rerank: service unavailable: {why}");
162                Err(RankError(why))
163            }
164        }
165    }
166}

170pub struct RemoteIcons {
171    base_url: String,
172}

174impl RemoteIcons {
175    pub fn new(base_url: impl Into<String>) -> Self {
176        Self { base_url: base_url.into() }
177    }
178}
179
180impl IconChooser for RemoteIcons {
181    fn choose(&self, text: &str) -> Result<Option<IconId>, IconError> {
182        debug!("remote icon choice for a fact of {} chars", text.len());
183        let body = serde_json::to_string(&IconRequest { text: text.to_owned() }).map_err(|e| IconError(e.to_string()))?;
184        let reply = post(&self.base_url, "/icon", body, 30).map_err(IconError)?;
185        match parse::<IconReply>("icon", &reply).map_err(IconError)? {
186            IconReply::Icon(icon) => Ok(Some(icon)),
187            IconReply::NoneFits => Ok(None),
188            IconReply::Unavailable(why) => {
189                warn!("icon: service unavailable: {why}");
190                Err(IconError(why))
191            }
192        }
193    }
194}

198pub struct RemoteSync {
199    base_url: String,
200}

202impl RemoteSync {
203    pub fn new(base_url: impl Into<String>) -> Self {
204        trace!("remote sync client");
205        Self { base_url: base_url.into() }
206    }
207
208    pub fn memory(&self, mine: &MemorySnapshot) -> Result<MemorySnapshot, String> {
209        info!("sync memory: sending {} facts, {} forgotten", mine.facts.len(), mine.forgotten.len());
210        let body = serde_json::to_string(mine).map_err(|e| e.to_string())?;
211        let theirs: MemorySnapshot = parse("memory sync", &post(&self.base_url, "/memory/sync", body, 15)?)?;
212        info!("sync memory: received {} facts, {} forgotten", theirs.facts.len(), theirs.forgotten.len());
213        Ok(theirs)
214    }
215
216    pub fn household(&self, mine: &Household) -> Result<Household, String> {
217        info!("sync household: sending the document");
218        let body = serde_json::to_string(mine).map_err(|e| e.to_string())?;
219        parse("household sync", &post(&self.base_url, "/household/sync", body, 10)?)
220    }
221}
222
223impl RemoteSync {

225    pub fn pictures_missing(&self, ids: &[String]) -> Result<Vec<String>, String> {
226        debug!("sync pictures: asking which of {} are missing", ids.len());
227        let body = serde_json::json!({ "ids": ids }).to_string();
228        let reply: serde_json::Value = parse("picture missing", &post(&self.base_url, "/picture/missing", body, 10)?)?;
229        let missing: Vec<String> =
230            reply["missing"].as_array().map(|a| a.iter().filter_map(|v| v.as_str().map(str::to_owned)).collect()).unwrap_or_default();
231        debug!("sync pictures: {} missing on the service", missing.len());
232        Ok(missing)
233    }

235    pub fn put_picture(&self, id: &str, bytes: &[u8]) -> Result<(), String> {
236        use base64::Engine as _;
237        debug!("sync pictures: putting {id} ({} bytes)", bytes.len());
238        let body = serde_json::json!({ "id": id, "data_base64": base64::engine::general_purpose::STANDARD.encode(bytes) }).to_string();
239        post(&self.base_url, "/picture/put", body, 30).map(|_| ())
240    }

243    pub fn get_picture(&self, id: &str) -> Result<Option<Vec<u8>>, String> {
244        use base64::Engine as _;
245        debug!("sync pictures: getting {id}");
246        let body = serde_json::json!({ "id": id }).to_string();
247        let reply: serde_json::Value = parse("picture get", &post(&self.base_url, "/picture/get", body, 30)?)?;
248        match reply["data_base64"].as_str() {
249            Some(d) => base64::engine::general_purpose::STANDARD.decode(d).map(Some).map_err(|e| {
250                warn!("sync pictures: {id} arrived as bad base64: {e}");
251                e.to_string()
252            }),
253            None => {
254                debug!("sync pictures: the service does not have {id}");
255                Ok(None)
256            }
257        }
258    }
259}

261impl RemoteSync {
262    pub fn chat(&self, mine: &whiskers_core::ChatState) -> Result<whiskers_core::ChatState, String> {
263        info!("sync chat: sending version {}", mine.version);
264        let body = serde_json::to_string(mine).map_err(|e| e.to_string())?;
265        let theirs: whiskers_core::ChatState = parse("chat sync", &post(&self.base_url, "/chat/sync", body, 10)?)?;
266        info!("sync chat: received version {}", theirs.version);
267        Ok(theirs)
268    }

271    pub fn journal_have(&self, device: &str) -> Result<usize, String> {
272        debug!("sync journal: asking how many lines the service has from {device}");
273        let body = serde_json::json!({ "device": device, "from": 0, "lines": [] }).to_string();
274        let r: serde_json::Value = parse("journal have", &post(&self.base_url, "/journal/push", body, 10)?)?;
275        r["have"].as_u64().map(|n| n as usize).ok_or_else(|| {
276            warn!("sync journal: the reply has no count");
277            "no count in the reply".to_owned()
278        })
279    }

283    pub fn journal_push(&self, device: &str, from: usize, lines: &[String]) -> Result<usize, String> {
284        info!("sync journal: pushing {} line(s) from line {from} for {device}", lines.len());
285        let body = serde_json::json!({ "device": device, "from": from, "lines": lines }).to_string();
286        let r: serde_json::Value = parse("journal push", &post(&self.base_url, "/journal/push", body, 30)?)?;
287        let have = r["have"].as_u64().map(|n| n as usize).ok_or_else(|| {
288            warn!("sync journal: the reply has no count");
289            "no count in the reply".to_owned()
290        })?;
291        if have != from + lines.len() {
292            warn!("sync journal: the service has {have} lines from {device}, expected {}", from + lines.len());
293        }
294        Ok(have)
295    }

299    pub fn journal_pull(&self, have: usize) -> Result<Pulled, String> {
300        info!("sync journal: pulling from line {have}");
301        let body = serde_json::json!({ "have": have }).to_string();
302        let r: serde_json::Value = parse("journal pull", &post(&self.base_url, "/journal/pull", body, 30)?)?;
303        let pulled = Pulled {
304            from: r["from"].as_u64().unwrap_or(0) as usize,
305            total: r["total"].as_u64().unwrap_or(0) as usize,
306            lines: r["lines"].as_array().map(|a| a.iter().filter_map(|x| x.as_str().map(str::to_owned)).collect()).unwrap_or_default(),
307        };
308        info!("sync journal: pulled {} line(s) from {} of {}", pulled.lines.len(), pulled.from, pulled.total);
309        Ok(pulled)
310    }
311}

314pub struct Pulled {
315    pub from: usize,
316    pub total: usize,
317    pub lines: Vec<String>,
318}

321#[derive(Clone, Copy, Debug, PartialEq, Eq)]
322pub enum ServiceProbe {

324    Reachable,

326    NotWhiskers,

328    NotReachable,
329}

335pub fn probe_service(address: &whiskers_core::ServiceAddress) -> ServiceProbe {
336    let started = Instant::now();
337    let agent: ureq::Agent = ureq::Agent::config_builder()
338        .timeout_global(Some(Duration::from_secs(12)))
339        .timeout_connect(Some(Duration::from_secs(5)))
340        .http_status_as_error(false)
341        .build()
342        .into();
343    let sent = agent.post(format!("{}/usage", address.url())).header("content-type", "application/json").send("{}");
344    let mut resp = match sent {
345        Ok(r) => r,
346        Err(e) => {
347            info!("service probe: no answer after {} ms ({})", started.elapsed().as_millis(), error_kind(&e));
348            return ServiceProbe::NotReachable;
349        }
350    };
351    let status = resp.status();
352    let body = resp.body_mut().read_to_string().unwrap_or_default();
353    let looks_right = status.is_success() && serde_json::from_str::<serde_json::Value>(&body).is_ok_and(|v| v.get("voice").is_some_and(|x| x.is_object()));
354    info!("service probe: answered {status} in {} ms, {} bytes, whiskers = {looks_right}", started.elapsed().as_millis(), body.len());
355    if looks_right { ServiceProbe::Reachable } else { ServiceProbe::NotWhiskers }
356}

359fn error_kind(e: &ureq::Error) -> &'static str {
360    match e {
361        ureq::Error::Timeout(_) => "timeout",
362        ureq::Error::Io(_) | ureq::Error::ConnectionFailed => "connection failed",
363        ureq::Error::HostNotFound => "host not found",
364        _ => "other",
365    }
366}

368#[cfg(test)]
369mod probe_tests {
370    use super::*;
371    use std::io::{Read, Write};
372    use std::net::TcpListener;

375    fn serve_once(status: &str, body: &'static str) -> whiskers_core::ServiceAddress {
376        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
377        let port = listener.local_addr().unwrap().port();
378        let status = status.to_owned();
379        std::thread::spawn(move || {
380            if let Ok((mut s, _)) = listener.accept() {
381                // Read the whole request (headers, then the declared body) before answering, or closing
382                // the socket with bytes unread resets the connection under the client.
383                let mut got = Vec::new();
384                let mut chunk = [0u8; 1024];
385                loop {
386                    let n = s.read(&mut chunk).unwrap_or(0);
387                    got.extend_from_slice(&chunk[..n]);
388                    let text = String::from_utf8_lossy(&got).into_owned();
389                    if let Some(end) = text.find("\r\n\r\n") {
390                        let want = text[..end].lines().find_map(|l| l.to_ascii_lowercase().strip_prefix("content-length:").and_then(|v| v.trim().parse::<usize>().ok())).unwrap_or(0);
391                        if got.len() >= end + 4 + want {
392                            break;
393                        }
394                    }
395                    if n == 0 {
396                        break;
397                    }
398                }
399                let _ = write!(s, "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", body.len());
400            }
401        });
402        whiskers_core::ServiceAddress::parse(&format!("127.0.0.1:{port}")).unwrap()
403    }

405    #[test]
406    fn the_whiskers_service_is_reachable() {
407        let a = serve_once("200 OK", r#"{"voice":{"configured":false},"tokens":{"used":0}}"#);
408        assert_eq!(probe_service(&a), ServiceProbe::Reachable);
409    }
410
411    #[test]
412    fn another_program_is_not_whiskers() {
413        assert_eq!(probe_service(&serve_once("200 OK", "<html>hello</html>")), ServiceProbe::NotWhiskers);
414        assert_eq!(probe_service(&serve_once("404 Not Found", "{}")), ServiceProbe::NotWhiskers);
415        assert_eq!(probe_service(&serve_once("500 Internal Server Error", r#"{"voice":{}}"#)), ServiceProbe::NotWhiskers);
416    }
417
418    #[test]
419    fn nothing_listening_is_not_reachable() {
420        let port = TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port();
421        let a = whiskers_core::ServiceAddress::parse(&format!("127.0.0.1:{port}")).unwrap();
422        assert_eq!(probe_service(&a), ServiceProbe::NotReachable);
423    }
424}