lune/src/lib/globals/net.rs

use std::{collections::HashMap, str::FromStr};

use mlua::{Error, Lua, LuaSerdeExt, Result, Table, Value};
use reqwest::{
    header::{HeaderMap, HeaderName, HeaderValue},
    Method,
};

use crate::utils::{net::get_request_user_agent_header, table_builder::TableBuilder};

pub async fn create(lua: &Lua) -> Result<()> {
    lua.globals().raw_set(
        "net",
        TableBuilder::new(lua)?
            .with_function("jsonEncode", net_json_encode)?
            .with_function("jsonDecode", net_json_decode)?
            .with_async_function("request", net_request)?
            .build_readonly()?,
    )
}

fn net_json_encode(_: &Lua, (val, pretty): (Value, Option<bool>)) -> Result<String> {
    if let Some(true) = pretty {
        serde_json::to_string_pretty(&val).map_err(Error::external)
    } else {
        serde_json::to_string(&val).map_err(Error::external)
    }
}

fn net_json_decode(lua: &Lua, json: String) -> Result<Value> {
    let json: serde_json::Value = serde_json::from_str(&json).map_err(Error::external)?;
    lua.to_value(&json)
}

async fn net_request<'lua>(lua: &'lua Lua, config: Value<'lua>) -> Result<Table<'lua>> {
    // Extract stuff from config and make sure its all valid
    let (url, method, headers, body) = match config {
        Value::String(s) => {
            let url = s.to_string_lossy().to_string();
            let method = "GET".to_string();
            (url, method, None, None)
        }
        Value::Table(tab) => {
            // Extract url
            let url = match tab.raw_get::<&str, mlua::String>("url") {
                Ok(config_url) => config_url.to_string_lossy().to_string(),
                Err(_) => {
                    return Err(Error::RuntimeError(
                        "Missing 'url' in request config".to_string(),
                    ))
                }
            };
            // Extract method
            let method = match tab.raw_get::<&str, mlua::String>("method") {
                Ok(config_method) => config_method.to_string_lossy().trim().to_ascii_uppercase(),
                Err(_) => "GET".to_string(),
            };
            // Extract headers
            let headers = match tab.raw_get::<&str, mlua::Table>("headers") {
                Ok(config_headers) => {
                    let mut lua_headers = HeaderMap::new();
                    for pair in config_headers.pairs::<mlua::String, mlua::String>() {
                        let (key, value) = pair?;
                        lua_headers.insert(
                            HeaderName::from_str(key.to_str()?).map_err(Error::external)?,
                            HeaderValue::from_str(value.to_str()?).map_err(Error::external)?,
                        );
                    }
                    Some(lua_headers)
                }
                Err(_) => None,
            };
            // Extract body
            let body = match tab.raw_get::<&str, mlua::String>("body") {
                Ok(config_body) => Some(config_body.as_bytes().to_owned()),
                Err(_) => None,
            };
            (url, method, headers, body)
        }
        value => {
            return Err(Error::RuntimeError(format!(
                "Invalid request config - expected string or table, got {}",
                value.type_name()
            )))
        }
    };
    // Convert method string into proper enum
    let method = match Method::from_str(&method) {
        Ok(meth) => meth,
        Err(_) => {
            return Err(Error::RuntimeError(format!(
                "Invalid request config method '{}'",
                &method
            )))
        }
    };
    // Extract headers from config, force user agent
    let mut header_map = if let Some(headers) = headers {
        headers
    } else {
        HeaderMap::new()
    };
    header_map.insert(
        "User-Agent",
        HeaderValue::from_str(&get_request_user_agent_header()).map_err(Error::external)?,
    );
    // Create a client to send a request with
    // FUTURE: Try to reuse this client
    let client = reqwest::Client::builder()
        .build()
        .map_err(Error::external)?;
    // Create and send the request
    let mut request = client.request(method, url).headers(header_map);
    if let Some(body) = body {
        request = request.body(body);
    }
    let response = request.send().await.map_err(Error::external)?;
    // Extract status, headers, body
    let res_status = response.status();
    let res_headers = response.headers().clone();
    let res_bytes = response.bytes().await.map_err(Error::external)?;
    // Construct and return a readonly lua table with results
    TableBuilder::new(lua)?
        .with_value("ok", res_status.is_success())?
        .with_value("statusCode", res_status.as_u16())?
        .with_value(
            "statusMessage",
            res_status.canonical_reason().unwrap_or("?"),
        )?
        .with_value(
            "headers",
            res_headers
                .iter()
                .filter_map(|(key, value)| match value.to_str() {
                    Ok(value) => Some((key.as_str(), value)),
                    Err(_) => None,
                })
                .collect::<HashMap<_, _>>(),
        )?
        .with_value("body", lua.create_string(&res_bytes)?)?
        .build_readonly()
}