leetbot/src/main.rs

mod lib;
use crate::lib::*;

use chrono::{NaiveDateTime, Local};
use sqlx::AnyPool;
use sqlx::any::AnyPoolOptions;
use teloxide::{prelude::*, RequestError, dispatching::UpdateFilterExt, utils::command::BotCommands};
use thiserror::Error;

// DB
async fn init_db_pool(url: &str) -> Result<AnyPool, sqlx::Error> {
    let pool = AnyPoolOptions::new()
        .max_connections(5)
        .connect(url).await?;
    
    Ok(pool)
}

async fn setup_tables(pool: &sqlx::AnyPool) -> Result<(), sqlx::Error> {
    
    let query_sql = match pool.any_kind() {
        sqlx::any::AnyKind::Postgres => "create table if not exists leet_log (
id int generated always as identity,
username varchar(255) not null,
log_time timestamp not null
)" ,
        sqlx::any::AnyKind::Sqlite => "create table if not exists leet_log (
id integer primary key autoincrement not null,
username varchar(255) not null,
log_time timestamp not null
)"
    };
    
    sqlx::query(query_sql)
        .execute(pool)
        .await?;

    Ok(())
}

#[derive(Debug,Error)]
enum LeetError {
    #[error("It's not leet right now")]
    NotLeet,
    #[error("Already hit leet for today")]
    AlreadyHitLeet,
    #[error("Sql error while checking for leet status: {0}")]
    DbError(#[from] sqlx::Error)
}

async fn check_leet(username: &str, pool: &AnyPool) -> Result<(), LeetError> {
    if ! time_is_leet() {
	return Err(LeetError::NotLeet);
    }

    let now = Local::now().naive_local();

    let latest_leet : Option<(i32, String, NaiveDateTime)> = sqlx::query_as("select * from leet_log where username = $1 order by log_time desc")
        .bind(username)
        .fetch_optional(pool)
        .await?;

    if let Some(latest_leet) = latest_leet {
	let distance = now - latest_leet.2;

	if distance.num_hours() < 24 {
	    return Err(LeetError::AlreadyHitLeet);
	}
    }

    Ok(())
}

async fn log_leet(username: &str, pool: &AnyPool) -> Result<(), sqlx::Error> {
    let sql_now = match pool.any_kind() {
        sqlx::any::AnyKind::Postgres => "now()",
        sqlx::any::AnyKind::Sqlite => "datetime()"
    };

    let query_sql = format!("insert into leet_log (username, log_time) values ($1, {})", sql_now);
    
    sqlx::query(&query_sql)
        .bind(username)
        .execute(pool)
        .await?;

    Ok(())
}

async fn get_scores(pool: &AnyPool) -> Result<Vec<(String,i64)>, sqlx::Error> {
    let score: Vec<_> = sqlx::query_as("select username, count(username) as score from leet_log group by username")
        .fetch_all(pool)
        .await?;

    Ok(score)
}


// Bot

fn message_ok() -> Result<(), MessageHandlerError> {
    Ok(respond(())?)
}


#[derive(BotCommands, Clone)]
#[command(rename = "lowercase", description = "score commands")]
enum LeetCommand {
    #[command(description = "get scores")]
    Scores
}

#[derive(Error,Debug)]
enum MessageHandlerError {
    #[error("Telegram request error {0}")]
    Request(#[from] RequestError),
    #[error("No sender for message!")]
    NoFrom,
    #[error("Sender has no username!")]
    NoUsername,
    #[error("Internal sql error {0}")]
    Db(#[from] sqlx::Error),
    #[error("No score for user in db")]
    NoScore,
}

fn filter_leet(message: Message) -> bool {
    let text = match message.text() {
	Some(t) => t,
	None => { return false; }
    };

    text_is_leet(text)
}

async fn leet_message_handler(message: Message, bot: AutoSend<Bot>, pool: AnyPool) -> Result<(), MessageHandlerError> {

    let username = message.from()
	.ok_or(MessageHandlerError::NoFrom)?
	.username.clone()
	.ok_or(MessageHandlerError::NoUsername)?;

    let leet_check = check_leet(&username, &pool).await;

    if let Err(e) = leet_check {
	match e {
	    LeetError::NotLeet => bot.send_message(message.chat.id, "it is not leet right now").await?,
	    LeetError::AlreadyHitLeet => bot.send_message(message.chat.id, "you already hit leet today!").await?,
	    LeetError::DbError(e) => {
		bot.send_message(message.chat.id, "internal sql error").await?;
		return Err(MessageHandlerError::Db(e));
	    }
	};
	return message_ok();
    }

    log_leet(&username, &pool).await?;

    let scores = get_scores(&pool).await?;

    let user_score = scores.iter().find(|(u,_)| u == &username).ok_or(MessageHandlerError::NoScore)?.1;

    let msg = format!("{} just hit leet! New score: {}", username, user_score);

    bot.send_message(message.chat.id, msg).await?;

    message_ok()
}

async fn leet_command_handler(message: Message, cmd: LeetCommand, bot: AutoSend<Bot>, pool: AnyPool) -> Result<(), MessageHandlerError> {

    match cmd {
        LeetCommand::Scores => {
	    let scores = get_scores(&pool).await?;

	    let scores_str = scores.into_iter().map(|(name,score)| format!(" {}: {}\n", name, score)).collect::<String>();

	    let msg = format!("Leet scores:\n{}", scores_str);

	    bot.send_message(message.chat.id, msg).await?;
	}
    }

    message_ok()
}


#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    pretty_env_logger::init();

    let db_url = std::env::var("DB_URL")
        .unwrap_or_else(|_| "sqlite::memory:".to_string());

    let db_pool = init_db_pool(&db_url).await?;

    setup_tables(&db_pool).await?;

    let telegram_token = std::env::var("TELEGRAM_BOT_TOKEN")
        .expect("TELEGRAM_TOKEN env not set!");

    let bot = Bot::new(&telegram_token).auto_send();

    let handler = Update::filter_message()
        .branch(dptree::entry()
		.filter_command::<LeetCommand>()
		.endpoint(leet_command_handler))
        .branch(dptree::entry()
		.filter(filter_leet)
		.endpoint(leet_message_handler))
	;
    
    Dispatcher::builder(bot, handler)
        .dependencies(dptree::deps![db_pool])
        .default_handler(|upd| async move {
	    println!("unhandled update: {:?}", upd);
	})
        .error_handler(LoggingErrorHandler::with_custom_text("error during dispatch"))
        .enable_ctrlc_handler()
        .build()
        .dispatch()
        .await;

    Ok(())
}