telegram-bot2 0.3.7

use crate::__private::{GenericHandler, HandlerError};
use crate::models::{GetUpdates, Update, UpdateType};
use crate::{Bot, BotState};
use std::env::var;
use std::fs::File;
use std::io::Read;

#[cfg(feature = "commands")]
use {crate::__private::command::CommandHandler, crate::models::MessageEntityType, lazy_static::lazy_static, regex::Regex, std::collections::BTreeMap};
#[cfg(feature = "daemons")]
use {crate::__private::daemon::DaemonStruct, tokio::sync::Mutex};

use log::{error, info, trace};
use std::sync::Arc;
use std::time::Duration;

/// Builds a bot.
/// This should be used in combination with the [bot][macro@crate::bot] macro
pub struct BotBuilder<T>
where
    T: Sync + Send + 'static,
{
    bot: Arc<Bot>,
    state: Arc<Option<BotState<T>>>,

    #[cfg(feature = "daemons")]
    daemons: Vec<DaemonStruct<T>>,

    #[cfg(feature = "commands")]
    commands: Option<BTreeMap<String, Vec<CommandHandler<T>>>>,

    handler: Vec<GenericHandler<T>>,

    update_limit: Option<i64>,
    allowed_updates: Option<Vec<UpdateType>>,
    offset: Option<i64>,
    timeout: Option<i64>,
    interval: Option<Duration>,
}

impl BotBuilder<()> {
    /// Builds a bot and uses the environment variable `BOT_TOKEN` as token or extract it from the `BOT_TOKEN_FILE` file. If both are defined, `BOT_TOKEN` is used.
    pub fn new() -> BotBuilder<()> {
        BotBuilder {
            bot: Arc::new(Bot::login(
                var("BOT_TOKEN")
                    .ok()
                    .or(var("BOT_TOKEN_FILE").ok().and_then(|f| {
                        let mut s = String::new();
                        File::open(f).expect("Could not open token file").read_to_string(&mut s).expect("Could not read from token file");
                        Some(s)
                    }))
                    .expect("No token provided (define BOT_TOKEN or BOT_TOKEN_FILE variable)"),
            )),
            state: Arc::new(None),
            #[cfg(feature = "commands")]
            commands: None,
            #[cfg(feature = "daemons")]
            daemons: vec![],
            handler: vec![],
            update_limit: None,
            allowed_updates: None,
            offset: None,
            timeout: None,
            interval: None,
        }
    }

    /// Builds a bot with the given token
    pub fn new_with_token(token: String) -> BotBuilder<()> {
        BotBuilder {
            bot: Arc::new(Bot::login(token)),
            state: Arc::new(None),
            #[cfg(feature = "commands")]
            commands: None,
            #[cfg(feature = "daemons")]
            daemons: vec![],
            handler: vec![],
            update_limit: None,
            allowed_updates: None,
            offset: None,
            timeout: None,
            interval: None,
        }
    }
}

impl Default for BotBuilder<()> {
    fn default() -> Self {
        BotBuilder::new()
    }
}

impl<T: Sync + Send + 'static> BotBuilder<T> {
    /// Add a shared state to the bot
    /// The state can be used by any update or command handler
    pub fn with_state<S>(self, state: S) -> BotBuilder<S>
    where
        S: Sync + Send + 'static,
    {
        if self.state.is_some() {
            panic!("State already set")
        }
        if !self.handler.is_empty() {
            panic!("Cannot set state after handler")
        }
        #[cfg(feature = "daemons")]
        if !self.daemons.is_empty() {
            panic!("Cannot set state after daemons")
        }
        #[cfg(feature = "commands")]
        if self.commands.is_some() {
            panic!("Cannot set state after commands")
        }

        BotBuilder {
            bot: self.bot,
            state: Arc::new(Some(BotState { state })),
            #[cfg(feature = "commands")]
            commands: None,
            #[cfg(feature = "daemons")]
            daemons: vec![],
            handler: vec![],
            update_limit: self.update_limit,
            allowed_updates: self.allowed_updates,
            offset: self.offset,
            timeout: self.timeout,
            interval: self.interval,
        }
    }

    /// Set the update handler of the bot
    pub fn handlers(mut self, handlers: Vec<GenericHandler<T>>) -> Self {
        if !self.handler.is_empty() {
            panic!("Handlers already set")
        }

        handlers.iter().for_each(|h| info!("Registered handler {} with priority {}", h.name, h.rank));

        self.handler = handlers;
        self.handler.sort_by_key(|f| f.rank);
        self
    }

    /// Set the maximum fetched updates per poll
    /// Default is 100, and max is 100
    pub fn update_limit(mut self, limit: i64) -> BotBuilder<T> {
        if self.update_limit.is_some() {
            panic!("Update limit already set")
        }

        self.update_limit = Some(limit);
        self
    }

    /// Specify the time (in seconds) before the polling exits with timeout
    /// Should be a stricly positive integer
    pub fn timeout(mut self, timeout: i64) -> BotBuilder<T> {
        if self.timeout.is_some() {
            panic!("Timeout already set")
        }
        // TODO warn for short timeouts

        self.timeout = Some(timeout);
        self
    }

    /// Minimum waiting time between two pollings
    pub fn interval(mut self, interval: Duration) -> BotBuilder<T> {
        if self.interval.is_some() {
            panic!("Interval already set")
        }

        self.interval = Some(interval);
        self
    }

    #[cfg(feature = "commands")]
    /// Set the command handlers (see [`commands`][crate::commands])
    pub fn commands(mut self, commands: BTreeMap<String, Vec<CommandHandler<T>>>) -> BotBuilder<T> {
        if self.commands.is_some() {
            panic!("Commands already set")
        }

        for cmd in &commands {
            info!("Registering handler for command {}: ", cmd.0);
            for handler in cmd.1 {
                info!(r#" - "{}" with priority {}"#, handler.syntax, handler.rank);
            }
        }

        self.commands = Some(commands);
        self
    }

    #[cfg(feature = "daemons")]
    /// Set the daemons of the bot (see [`daemon`][crate::daemon])
    pub fn daemons(mut self, daemons: Vec<DaemonStruct<T>>) -> BotBuilder<T> {
        if !self.daemons.is_empty() {
            panic!("Daemons already set")
        }

        daemons.iter().for_each(|d| info!("Registering daemon {}", d.syntax));

        self.daemons = daemons;
        self
    }

    /// Start the bot (should not be used, see [`bot`][macro@crate::bot])
    pub async fn launch(mut self) {
        let allowed_updates = self.allowed_updates.unwrap_or_default();
        let timeout = self.timeout.or(Some(5));

        info!("Launching bot with parameters: [interval: {:?}, timeout: {:?}]", self.interval, timeout);

        #[cfg(feature = "daemons")]
        Self::launch_daemons(self.daemons, self.bot.clone(), self.state.clone()).await;

        loop {
            if let Some(interval) = self.interval {
                tokio::time::sleep(interval).await;
            }

            // Wait for an update to come
            let mut updates = vec![];
            while updates.is_empty() {
                updates = self
                    .bot
                    .get_updates(GetUpdates {
                        offset: self.offset,
                        limit: self.update_limit,
                        timeout,
                        allowed_updates: {
                            #[cfg(feature = "commands")]
                            if self.commands.is_some() {
                                [allowed_updates.clone(), vec![UpdateType::NewMessage]].concat()
                            } else {
                                allowed_updates.clone()
                            }

                            #[cfg(not(feature = "commands"))]
                            allowed_updates.clone()
                        },
                    })
                    .await
                    .unwrap()
            }

            // Sort all updates by their id such that they are processed correctly
            updates.sort_by_key(|u| u.update_id);

            // Set the offset to the id after the last one recieved
            // We may unwrap it, as the loop above ensure that the vec is not empty
            self.offset = Some(updates.last().unwrap().update_id + 1);

            for update in updates {
                // Check whether the update is a command and there is a valid registered handleer
                if let Ok(update) = Update::try_from(update.clone()) {
                    trace!("Received update {} ({:?})", update.get_id(), update.get_type());

                    #[cfg(feature = "commands")]
                    if Self::try_commands(&self.commands, &self.bot, self.state.as_ref().as_ref(), &update).await {
                        continue;
                    }

                    // Else pass the update to the default handler
                    for h in &self.handler {
                        if h.updates.is_empty() || h.updates.iter().any(|u| *u == update.get_type()) {
                            match h.handler.handle(&self.bot, &update, self.state.as_ref().as_ref()).await {
                                Ok(_) => break,
                                Err(HandlerError::Parse) => continue,
                                Err(HandlerError::Runtime) => {
                                    error!("Handler failure for {:?}", update);
                                    break;
                                }
                            }
                        }
                    }
                } else {
                    error!("Invalid update received: {:?}", update)
                }
            }
        }
    }

    #[cfg(feature = "commands")]
    /// Tries to parse the update for as a command. Returns whether it was successful
    async fn try_commands(commands: &Option<BTreeMap<String, Vec<CommandHandler<T>>>>, bot: &Bot, state: Option<&BotState<T>>, update: &Update) -> bool {
        if let (Update::NewMessage(_, m), Some(commands)) = (&update, commands) {
            if let Some(e) = m.entities.iter().find(|e| e.offset == 0 && matches!(e._type, MessageEntityType::BotCommand)) {
                trace!("Command \"{}\" issued", m.text.as_ref().unwrap_or(&"".to_string()));

                lazy_static! {
                    static ref PARAMETER_DELIMITER: Regex = Regex::new(r"\s+").unwrap();
                }

                if let Some(handlers) = commands.get(&m.text.as_ref().unwrap().as_str()[(e.offset + 1) as usize..(e.offset + e.length) as usize]) {
                    let args: Vec<String> = PARAMETER_DELIMITER
                        .split(&m.text.as_ref().unwrap().as_str()[(e.offset + e.length) as usize..])
                        .filter(|s| !s.is_empty())
                        .map(str::to_string)
                        .collect();

                    for h in handlers {
                        match h.handler.handle(&args, update, bot, state).await {
                            Ok(_) => break,
                            Err(HandlerError::Parse) => continue,
                            Err(HandlerError::Runtime) => panic!(),
                        }
                    }

                    trace!("Could not delegate to any handler");

                    return true;
                } else {
                    trace!("No corresponding handler found")
                }
            }
        }

        false
    }

    #[cfg(feature = "daemons")]
    /// Launch all the bot daemons and returns their lock
    async fn launch_daemons(daemons: Vec<DaemonStruct<T>>, bot: Arc<Bot>, state: Arc<Option<BotState<T>>>) -> Vec<Arc<Mutex<()>>> {
        daemons
            .into_iter()
            .map(|d| {
                let mutex = Arc::new(Mutex::new(()));

                let bot = bot.clone();
                let state = state.clone();
                let lock = mutex.clone();

                tokio::task::spawn(async move { d.daemon.launch(bot, state, lock).await });

                mutex
            })
            .collect()
    }
}