path: root/src/commands/discocaml.rs

use anyhow::{anyhow, bail, Context as _, Error, Result};
use bstr::BString;
use futures::Future;
use serde::{Deserialize, Serialize};
use serde_json::de::Deserializer;
use serenity::{
    all::{
        CommandDataOptionValue, CommandInteraction, CommandOptionType, CommandType,
        ComponentInteraction, ComponentInteractionDataKind, InteractionResponseFlags, Member,
        RoleId,
    },
    builder::{
        CreateAttachment, CreateButton, CreateCommand, CreateCommandOption, CreateEmbed,
        CreateInteractionResponse, CreateInteractionResponseMessage,
    },
    client::Context,
    model::Permissions,
};
use sqlx::SqlitePool;
use std::{collections::BTreeSet, process::Stdio};
use tokio::{io::AsyncWriteExt, process::Command};

use crate::utils::EnumAsArray;

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct DiscocamlConfig {
    pub command: Vec<String>,
    pub roles: BTreeSet<RoleId>,
}

#[derive(Debug, PartialEq, Serialize)]
pub struct DiscocamlRequest {
    pub expr: String,
    pub command: DiscocamlCommand,
}

#[derive(Debug, PartialEq, Serialize)]
pub enum DiscocamlCommand {
    Parse,
    DrawTree,
    RunCBN,
    RunCBV,
    StepCBN,
    StepCBV,
}

/// A response outputted by the discocaml subprocess as a JSON string.
///
/// ```
/// # use lambo::{commands::discocaml::*, utils::EnumAsArray};
/// # use serde::Deserialize;
/// # use serde_json::Deserializer;
///
/// let example = r#"
/// [ "Expr"
/// , { "expr": "1 + 2"
///   , "has_redex": true
///   , "note": null
///   }
/// ]
/// "#;
/// let expected = DiscocamlResponse::Expr(DiscocamlResponseExpr {
///     expr: "1 + 2".to_string(),
///     has_redex: true,
///     note: None,
/// });
///
/// let mut de = Deserializer::from_str(&example);
/// let out = DiscocamlResponse::deserialize(EnumAsArray(&mut de)).unwrap();
/// de.end().unwrap();
///
/// assert_eq!(out, expected);
/// ```
#[derive(Debug, Deserialize, PartialEq)]
#[serde(deny_unknown_fields)]
pub enum DiscocamlResponse {
    Expr(DiscocamlResponseExpr),
    Error(String),
    Graphviz(String),
}

#[derive(Debug, Deserialize, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct DiscocamlResponseExpr {
    pub expr: String,
    pub has_redex: bool,
    pub note: Option<String>,
}

impl DiscocamlResponseExpr {
    async fn save(&self, db: &SqlitePool) -> Result<i64> {
        let out = sqlx::query!(
            "INSERT INTO discocaml_exprs (expr) VALUES (?) RETURNING rowid",
            self.expr
        )
        .fetch_one(db)
        .await
        .context("failed to save expression to database")?;
        Ok(out.rowid)
    }
}

pub fn command() -> CreateCommand {
    CreateCommand::new("discocaml")
        .kind(CommandType::ChatInput)
        .default_member_permissions(Permissions::empty())
        .dm_permission(true)
        .description("Sends this expression to the disco!")
        .add_option(
            CreateCommandOption::new(
                CommandOptionType::String,
                "expr",
                "The expression to operate on.",
            )
            .required(true),
        )
}

async fn check_permissions(config: &DiscocamlConfig, member: Option<&Member>) -> Result<()> {
    if let Some(member) = member {
        if !member.roles.iter().any(|role| config.roles.contains(role)) {
            match config.roles.len() {
                0 => bail!("This command can not be used, because roles are not configured."),
                1 => bail!(
                    "This command can only be used by <@&{}>.",
                    config.roles.iter().next().unwrap()
                ),
                2 => {
                    let mut iter = config.roles.iter();
                    bail!(
                        "This command can only be used by <@&{}> or <@&{}>.",
                        iter.next().unwrap(),
                        iter.next().unwrap()
                    )
                }
                _ => {
                    let mut iter = config.roles.iter();
                    let last = iter.next().unwrap();
                    let mut buf = String::new();
                    for role in iter {
                        buf.push_str("<@&");
                        buf.push_str(&role.to_string());
                        buf.push_str(">, ");
                    }
                    buf.push_str(", or <@&");
                    buf.push_str(&last.to_string());
                    buf.push('>');
                    bail!("This command can only be used by {}.", buf)
                }
            }
        }
        Ok(())
    } else {
        bail!("This command cannot be used in DMs.")
    }
}

async fn respond_with_error<E: std::error::Error, F: Future<Output = Result<(), E>>>(
    err: &Error,
    send: impl FnOnce(CreateInteractionResponse) -> F,
) {
    let msg = CreateInteractionResponseMessage::new()
        .content(format!(":no_entry_sign: {}", err))
        .flags(InteractionResponseFlags::EPHEMERAL);
    if let Err(err) = send(CreateInteractionResponse::Message(msg)).await {
        log::error!(
            "failed to respond to command that failed permissions check: {:?}",
            err
        )
    }
}

async fn respond_with<E, F>(
    db: &SqlitePool,
    res: DiscocamlResponse,
    send: impl FnOnce(CreateInteractionResponse) -> F,
) -> Result<()>
where
    E: std::error::Error + 'static + Send + Sync,
    F: Future<Output = Result<(), E>>,
{
    match res {
        DiscocamlResponse::Expr(expr) => {
            // Insert the output expression in the database.
            let rowid = match expr.save(db).await {
                Ok(rowid) => rowid,
                Err(err) => {
                    respond_with_error(&err, send).await;
                    return Err(err);
                }
            };

            // Respond with the expression and the buttons.
            let res = expr_response_message(&expr, rowid);
            send(CreateInteractionResponse::Message(res)).await?;
            Ok(())
        }
        DiscocamlResponse::Graphviz(dot) => {
            let png = match run_dot(&dot).await {
                Ok(png) => png,
                Err(err) => {
                    respond_with_error(&err, send).await;
                    return Err(err);
                }
            };

            let msg = CreateInteractionResponseMessage::new()
                .add_file(CreateAttachment::bytes(png, "expr.png"))
                .add_embed(CreateEmbed::new().attachment("expr.png"));
            send(CreateInteractionResponse::Message(msg)).await?;
            Ok(())
        }
        DiscocamlResponse::Error(err) => {
            let err = anyhow!("got an error from discocaml: `{}`", err);
            respond_with_error(&err, send).await;
            Err(err)
        }
    }
}

pub async fn run_discocaml(
    config: &DiscocamlConfig,
    req: &DiscocamlRequest,
) -> Result<DiscocamlResponse> {
    let mut child = Command::new(&config.command[0])
        .args(&config.command)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .kill_on_drop(true)
        .spawn()
        .context("failed to start discocaml")?;

    let mut stdin = child.stdin.take().unwrap();
    let req = serde_json::to_string(req).context("failed to serialize request to discocaml")?;
    stdin
        .write_all(req.as_bytes())
        .await
        .context("failed to write request to discocaml")?;
    stdin
        .shutdown()
        .await
        .context("failed to close pipe to discocaml")?;
    drop(stdin);
    drop(req);

    let output = child
        .wait_with_output()
        .await
        .context("failed to wait for discocaml to complete")?;
    if !output.status.success() {
        bail!("discocaml exited with non-zero status {:?}", output.status)
    }

    let mut de = Deserializer::from_slice(&output.stdout);
    let out = DiscocamlResponse::deserialize(EnumAsArray(&mut de))
        .context("failed to parse response from discocaml")?;
    de.end()
        .context("failed to parse response from discocaml")?;
    Ok(out)
}

async fn run_dot(dot: &str) -> Result<BString> {
    let mut child = Command::new("dot")
        .arg("-Tpng")
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .kill_on_drop(true)
        .spawn()
        .context("failed to start dot")?;

    let mut stdin = child.stdin.take().unwrap();
    stdin
        .write_all(dot.as_bytes())
        .await
        .context("failed to write request to dot")?;
    stdin
        .shutdown()
        .await
        .context("failed to close pipe to dot")?;
    drop(stdin);

    let output = child
        .wait_with_output()
        .await
        .context("failed to wait for dot to complete")?;
    if !output.status.success() {
        bail!("dot exited with non-zero status {:?}", output.status)
    }
    Ok(BString::new(output.stdout))
}

fn expr_response_message(
    expr: &DiscocamlResponseExpr,
    rowid: i64,
) -> CreateInteractionResponseMessage {
    let mut msg = String::new();
    if let Some(note) = &expr.note {
        msg.push_str(note);
        msg.push('\n');
    }
    msg.push_str("```ocaml\n");
    // TODO: Escaping
    msg.push_str(&expr.expr);
    msg.push_str("```\n");

    CreateInteractionResponseMessage::new()
        .content(msg)
        .button(CreateButton::new(format!("discocaml-draw-tree/{}", rowid)).label("Draw Tree"))
        .button(
            CreateButton::new(format!("discocaml-step-cbv/{}", rowid))
                .label("Step (CBV)")
                .disabled(!expr.has_redex),
        )
        .button(
            CreateButton::new(format!("discocaml-step-cbn/{}", rowid))
                .label("Step (CBN)")
                .disabled(!expr.has_redex),
        )
        .button(
            CreateButton::new(format!("discocaml-run-cbv/{}", rowid))
                .label("Run (CBV)")
                .disabled(!expr.has_redex),
        )
        .button(
            CreateButton::new(format!("discocaml-run-cbn/{}", rowid))
                .label("Run (CBN)")
                .disabled(!expr.has_redex),
        )
}

pub async fn handle_command(
    ctx: &Context,
    config: &DiscocamlConfig,
    db: &SqlitePool,
    command: &CommandInteraction,
) -> Result<()> {
    // Check that the required role was present.
    if let Err(err) = check_permissions(config, command.member.as_deref()).await {
        respond_with_error(&err, |res| command.create_response(&ctx, res)).await;
        return Err(err.context("permissions check failed"));
    }

    // Parse the expression out.
    let mut expr = None;
    for option in &command.data.options {
        match (&option.name as &str, &option.value) {
            ("expr", CommandDataOptionValue::String(value)) => expr = Some(value),
            _ => {
                let err = anyhow!("unknown option {:?}", option);
                respond_with_error(&err, |res| command.create_response(&ctx, res)).await;
                return Err(err);
            }
        }
    }
    let expr = if let Some(expr) = expr {
        expr
    } else {
        let err = anyhow!("missing option {:?}", "expr");
        respond_with_error(&err, |res| command.create_response(&ctx, res)).await;
        return Err(err);
    };

    // Round-trip the expression through discocaml.
    let req = DiscocamlRequest {
        expr: expr.to_string(),
        command: DiscocamlCommand::Parse,
    };
    let res = match run_discocaml(config, &req).await {
        Ok(res) => res,
        Err(err) => {
            respond_with_error(&err, |res| command.create_response(&ctx, res)).await;
            return Err(err.context("failed to run discocaml"));
        }
    };

    // Respond with the resulting expression.
    respond_with(db, res, |res| command.create_response(&ctx, res)).await
}

pub async fn handle_button(
    ctx: &Context,
    config: &DiscocamlConfig,
    db: &SqlitePool,
    component: &ComponentInteraction,
) -> Result<()> {
    // Check that the required role was present.
    if let Err(err) = check_permissions(config, component.member.as_ref()).await {
        respond_with_error(&err, |res| component.create_response(&ctx, res)).await;
        return Err(err.context("permissions check failed"));
    }

    // Come up with a command for discocaml based on which button was clicked.
    let (command, id) = match &component.data.kind {
        ComponentInteractionDataKind::Button => {
            let custom_id = &component.data.custom_id;
            if let Some(id_str) = custom_id.strip_prefix("discocaml-draw-tree/") {
                (DiscocamlCommand::DrawTree, id_str)
            } else if let Some(id_str) = custom_id.strip_prefix("discocaml-step-cbv/") {
                (DiscocamlCommand::StepCBV, id_str)
            } else if let Some(id_str) = custom_id.strip_prefix("discocaml-step-cbn/") {
                (DiscocamlCommand::StepCBN, id_str)
            } else if let Some(id_str) = custom_id.strip_prefix("discocaml-run-cbv/") {
                (DiscocamlCommand::RunCBV, id_str)
            } else if let Some(id_str) = custom_id.strip_prefix("discocaml-run-cbn/") {
                (DiscocamlCommand::RunCBN, id_str)
            } else {
                let err = anyhow!("unknown component {:?}", component.data);
                respond_with_error(&err, |res| component.create_response(&ctx, res)).await;
                return Err(err);
            }
        }
        _ => {
            let err = anyhow!("unknown component {:?}", component.data);
            respond_with_error(&err, |res| component.create_response(&ctx, res)).await;
            return Err(err);
        }
    };

    // Parse the ID, and fetch the corresponding expr.
    let id: i64 = id.parse().context("failed to parse id")?;
    let result = sqlx::query!("SELECT expr FROM discocaml_exprs WHERE rowid = ?", id)
        .fetch_one(db)
        .await
        .context("failed to load expression from database");
    let expr = match result {
        Ok(expr) => expr.expr,
        Err(err) => {
            respond_with_error(&err, |res| component.create_response(&ctx, res)).await;
            return Err(err);
        }
    };

    // Send discocaml the request.
    let req = DiscocamlRequest { expr, command };
    let res = match run_discocaml(config, &req).await {
        Ok(res) => res,
        Err(err) => {
            respond_with_error(&err, |res| component.create_response(&ctx, res)).await;
            return Err(err.context("failed to run discocaml"));
        }
    };

    // Respond.
    respond_with(db, res, |res| component.create_response(&ctx, res)).await
}