izzilis/src/selection/telegram.rs

use futures::{
    channel::oneshot::{self, Sender},
    future::BoxFuture,
    lock::Mutex,
    Sink, SinkExt, StreamExt,
};
use std::{collections::HashMap, fmt::Debug, sync::Arc};
use telegram_bot::{
    self, types::requests::answer_callback_query::CanAnswerCallbackQuery, Api, ChatRef,
    DeleteMessage, EditMessageReplyMarkup, InlineKeyboardButton, InlineKeyboardMarkup, Message,
    MessageKind, SendMessage, ToMessageId, ToSourceChat, Update, UpdateKind,
};
use uuid::Uuid;

use super::Selector;

#[derive(Clone)]
pub struct TelegramSelector {
    client: Arc<Api>,
    chat_ref: Arc<Mutex<ChatRef>>,
    pending: Arc<Mutex<HashMap<Uuid, Sender<bool>>>>,
}

impl TelegramSelector {
    pub fn new<S: Sink<ChatRef> + Send + Unpin + 'static>(
        api: Arc<Api>,
        chat_ref: Arc<Mutex<ChatRef>>,
        mut updates: S,
    ) -> Self
    where
        S::Error: Debug,
    {
        let pending: Arc<Mutex<HashMap<Uuid, Sender<bool>>>> = Arc::new(Mutex::new(HashMap::new()));
        let api_clone = api.clone();
        let chat_ref_clone = chat_ref.clone();
        let pending_clone = pending.clone();
        tokio::spawn(async move {
            let mut stream = api_clone.stream();
            while let Some(Ok(data)) = stream.next().await {
                if let Update {
                    kind: UpdateKind::CallbackQuery(query),
                    ..
                } = data
                {
                    if let Some(data) = query.data.clone() {
                        let uuid_bytes = data.as_bytes().get(0..32);
                        let bool_byte = data.as_bytes().get(32);
                        if let (Some(uuid), Some(keep)) = (uuid_bytes, bool_byte) {
                            let uuid = Uuid::parse_str(&String::from_utf8_lossy(uuid));
                            if let Ok(uuid) = uuid {
                                if let Some(sender) = pending_clone.lock().await.remove(&uuid) {
                                    let keep = match *keep as char {
                                        't' => true,
                                        _ => false,
                                    };
                                    let _ = sender.send(keep);
                                    let _ = api_clone
                                        .send(query.answer(if keep {
                                            "Kept!"
                                        } else {
                                            "Discarded!"
                                        }))
                                        .await;

                                    if let Some(message) = query.message {
                                        let _ = api_clone
                                            .send(EditMessageReplyMarkup::new(
                                                message.to_source_chat(),
                                                message.to_message_id(),
                                                None::<InlineKeyboardMarkup>,
                                            ))
                                            .await;
                                        if !keep {
                                            let _ = api_clone
                                                .send(DeleteMessage::new(
                                                    message.to_source_chat(),
                                                    message.to_message_id(),
                                                ))
                                                .await;
                                        }
                                    }

                                    continue;
                                }
                            }
                        }
                    }
                } else if let Update {
                    kind:
                        UpdateKind::Message(Message {
                            chat,
                            kind: MessageKind::Text { data, .. },
                            ..
                        }),
                    ..
                } = data
                {
                    if data.starts_with("/setmain") {
                        let new_chat_ref = ChatRef::from_chat_id(chat.id());
                        if let Err(e) = updates.send(new_chat_ref.clone()).await {
                            println!("failed to send updated chat ref: {:?}", e);
                        }
                        *chat_ref_clone.lock().await = new_chat_ref;
                    }
                }
            }
        });
        Self {
            client: api,
            chat_ref: chat_ref,
            pending: pending,
        }
    }
}

impl Selector for TelegramSelector {
    type Error = telegram_bot::Error;

    type Response = BoxFuture<'static, Result<bool, Self::Error>>;

    fn review(&self, message: String) -> Self::Response {
        let client = self.client.clone();
        let chat_ref = self.chat_ref.clone();
        let pending = self.pending.clone();
        Box::pin(async move {
            let chat_ref = chat_ref.lock().await.clone();
            let uuid = Uuid::new_v4();
            let generate_callback = |keep| {
                let mut buffer = vec![0u8; 33];
                uuid.to_simple().encode_lower(&mut buffer);
                buffer[32] = if keep { 't' } else { 'f' } as u8;
                String::from_utf8_lossy(&buffer).to_string()
            };
            let mut message = SendMessage::new(chat_ref, message);
            message.reply_markup({
                let mut kb = InlineKeyboardMarkup::new();
                kb.add_row(vec![
                    InlineKeyboardButton::callback("Keep", generate_callback(true)),
                    InlineKeyboardButton::callback("Discard", generate_callback(false)),
                ]);
                kb
            });
            let (sender, receiver) = oneshot::channel();
            pending.lock().await.insert(uuid, sender);
            client.send(message).await?;
            Ok(receiver.await.unwrap())
        })
    }
}

pub async fn get_chat_ref(
    api: Arc<Api>,
    mut chat_ref: Option<ChatRef>,
) -> Result<Arc<Mutex<ChatRef>>, telegram_bot::Error> {
    let mut stream = api.stream();

    if chat_ref.is_none() {
        while let Some(data) = stream.next().await.transpose()? {
            if let Update {
                kind:
                    UpdateKind::Message(Message {
                        chat,
                        kind: MessageKind::Text { data, .. },
                        ..
                    }),
                ..
            } = data
            {
                if data.starts_with("/setmain") {
                    let chat_ref_temp = ChatRef::from_chat_id(chat.id());
                    chat_ref = Some(chat_ref_temp);
                    break;
                }
            }
        }
    }

    let chat_ref = Arc::new(Mutex::new(chat_ref.expect("bot API failed silently")));

    Ok(chat_ref)
}