use crate::util::{quote_plus, truncate};
use anyhow::{bail, Context, Error, Result};
use figment::providers::Env;
use futures::join;
use irc::client::prelude::*;
use macros::privmsg;
use reqwest::{get, Url};
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct WolframAlpha {
wa_api_key: String,
}
impl WolframAlpha {
pub fn new(bot: &crate::Bot) -> Result<WolframAlpha> {
bot.figment
.clone()
.merge(Env::prefixed("CATINATOR_"))
.extract()
.context("failed to extract wolfram alpha config")
}
pub async fn wa(&self, bot: &crate::Bot, msg: Message) -> Result<()> {
privmsg!(msg, {
let content = get_input_query(text)?;
bot.send_privmsg(
msg.response_target()
.context("failed to get response target")?,
&wa_query(&content, Some(&self.wa_api_key), None).await?,
)?;
})
}
}
#[derive(Serialize, Deserialize, Debug)]
struct WaResponse {
queryresult: QueryResult,
}
#[derive(Serialize, Deserialize, Debug)]
struct QueryResult {
pods: Vec<Pod>,
}
#[derive(Serialize, Deserialize, Debug)]
struct Pod {
title: String,
id: String,
primary: Option<bool>,
subpods: Vec<SubPod>,
}
#[derive(Serialize, Deserialize, Debug)]
struct SubPod {
plaintext: String,
}
fn clean_result_text(text: &str) -> String {
text
.replace("\n", "; ")
.split_whitespace()
.collect::<Vec<&str>>()
.join(" ")
}
fn to_single_string(wa_res: WaResponse) -> String {
wa_res
.queryresult
.pods
.iter()
.filter(|it| it.id.to_lowercase() != "input" && it.primary.is_some())
.map(|pod| {
let subpod_texts = pod
.subpods
.iter()
.map(|subpod| clean_result_text(&subpod.plaintext))
.collect::<Vec<String>>()
.join(", ");
format!("{}: {}", &pod.title, subpod_texts)
})
.collect::<Vec<String>>()
.join(" - ")
}
fn get_wa_api_url(
query_str: &str,
api_key: Option<&str>,
base_url: Option<&str>,
) -> Result<Url, Error> {
let wa_url = "http://api.wolframalpha.com";
let api_url = format!(
"{}/v2/query?input={}&appid={}&output=json",
base_url.unwrap_or(wa_url),
quote_plus(query_str)?,
api_key.unwrap_or("XXX"), );
Url::parse(&api_url).context("Failed to parse URL")
}
async fn send_wa_req(url: &Url) -> Result<String, Error> {
let body = get(url.to_owned())
.await
.context("Failed to make request")?
.text()
.await
.context("failed to get request response text")?;
Ok(body)
}
async fn handle_wa_req(url: &Url) -> Result<WaResponse, Error> {
let res_body = send_wa_req(url).await?;
let parsed = serde_json::from_str(&res_body)?;
Ok(parsed)
}
async fn get_wa_user_short_url(input: &str) -> Result<String, Error> {
let user_url = format!(
"http://www.wolframalpha.com/input/?i={}",
quote_plus("e_plus(input)?)?
);
Ok(user_url)
}
#[tracing::instrument]
async fn wa_query(
query_str: &str,
api_key: Option<&str>,
base_url: Option<&str>,
) -> Result<String, Error> {
let user_url_shortened_fut = get_wa_user_short_url(query_str);
let url = get_wa_api_url(query_str, api_key, base_url)?;
let wa_res_fut = handle_wa_req(&url);
let futs = join!(wa_res_fut, user_url_shortened_fut);
let wa_res = match futs.0 {
Ok(x) => x,
_ => return Ok("No results.".to_string()),
};
let user_url_shortened = futs.1?;
let string_result = match to_single_string(wa_res) {
x if x.is_empty() => "No plaintext results.".to_string(),
x => x,
};
Ok(format!(
"{} - {}",
truncate(&string_result, 250), &user_url_shortened,
))
}
fn get_input_query(text: &str) -> Result<String, Error> {
let input = text.chars().as_str().splitn(2, " ").collect::<Vec<&str>>();
if input.len() != 2 {
bail!("Empty input for WA query");
}
let content = input[1].trim();
Ok(content.to_string())
}
#[cfg(test)]
mod tests {
use crate::hooks::wolfram_alpha::clean_result_text;
use super::{get_input_query, get_wa_user_short_url, wa_query};
use anyhow::{Error, Result};
use mockito::{self, Matcher};
#[test]
fn test_input_query_content_retrieval() -> Result<(), Error> {
let incoming = ":wa test";
let content = get_input_query(incoming)?;
assert_eq!(content, "test");
Ok(())
}
#[test]
fn test_input_query_content_retrieval_with_spaces() -> Result<(), Error> {
let incoming = ":wa foo bar";
let content = get_input_query(incoming)?;
assert_eq!(content, "foo bar");
Ok(())
}
#[test]
fn test_input_query_content_retrieval_with_more_spaces() -> Result<(), Error> {
let incoming = ":wa foo bar baz";
let content = get_input_query(incoming)?;
assert_eq!(content, "foo bar baz");
Ok(())
}
#[test]
fn test_clean_result_text() {
assert_eq!(
clean_result_text("Newlines\nand multiple\n\n whitespace is removed."),
"Newlines; and multiple; ; whitespace is removed.",
)
}
#[tokio::test]
async fn test_query_with_result_with_wrong_json_parsing() -> Result<(), Error> {
let body = include_str!("../../tests/resources/wolfram_alpha_api_response_wrong_json.json");
let _m = mockito::mock("GET", Matcher::Any)
.with_body(body)
.create();
mockito::start();
let res = wa_query("what is a url", None, Some(&mockito::server_url())).await?;
assert_eq!(res, "No results.");
Ok(())
}
}