pass_manager/src/state/get_user.rs

use crate::prelude::*;
use futures::TryFutureExt;
use itertools::Itertools;
use std::{borrow::Cow, fmt::Write, path::Path};
use teloxide::{
    net::Download,
    types::{Document, FileMeta},
};
use tokio::{task::spawn_blocking, try_join};
use trim_in_place::TrimInPlace;

#[inline]
fn validate_document(document: Option<&Document>) -> Result<&Document, &'static str> {
    let Some(document) = document else {
        return Err("You didn't send a file. Try again");
    };

    if document.file.size > 1024 * 1024 * 200 {
        return Err("The file is larger than 200 MiB. Try splitting it into multiple files");
    }

    let name = match document.file_name.as_deref() {
        Some(name) => Path::new(name.trim_end()),
        None => return Err("Couldn't get the name of the file. Try sending it again"),
    };

    match name.extension() {
        Some(ext) if ext.eq_ignore_ascii_case("json") => Ok(document),
        _ => Err("Invalid file name. You need to send a json file. Try again"),
    }
}

#[inline]
async fn download_file(bot: &Throttle<Bot>, file: &FileMeta) -> crate::Result<Vec<u8>> {
    let path = bot.get_file(file.id.as_str()).await?.path;
    let mut data = Vec::with_capacity(file.size as usize);
    bot.download_file_stream(&path)
        .try_for_each(|bytes| {
            data.extend_from_slice(&bytes);
            async { Ok(()) }
        })
        .await?;
    Ok(data)
}

#[inline]
fn process_accounts(
    accounts: &mut [DecryptedAccount],
    existing_names: ahash::HashSet<String>,
) -> crate::Result<Result<(), String>> {
    for account in accounts.iter_mut() {
        account.name.trim_in_place();
        account.login.trim_in_place();
        account.password.trim_in_place();
    }

    accounts.sort_unstable_by(|a, b| a.name.cmp(&b.name));

    let mut duplicates = Vec::new();
    let mut existing = Vec::new();
    let mut invalid = Vec::new();

    accounts
        .iter()
        .dedup_by_with_count(|a, b| a.name == b.name)
        .for_each(|(count, account)| {
            let duplicate = count != 1;
            let exists = existing_names.contains(&account.name);
            if duplicate {
                duplicates.push(account.name.as_str());
            }
            if exists {
                existing.push(account.name.as_str());
            }
            // If it already exists or if it is a duplicate there's no need to check the account's validity
            if !duplicate && !exists && !account.validate() {
                invalid.push(account.name.as_str());
            }
        });

    drop(existing_names);

    if duplicates.is_empty() && invalid.is_empty() && existing.is_empty() {
        return Ok(Ok(()));
    }

    let mut error_text = "Your accounts have the following problems:".to_owned();

    if !duplicates.is_empty() {
        write!(
            error_text,
            "\n\nDuplicate names:\n{:?}",
            duplicates.into_iter().format("\n")
        )?;
    }

    if !existing.is_empty() {
        write!(
            error_text,
            "\n\nAccounts with these names already exist in the database:\n{:?}",
            existing.into_iter().format("\n")
        )?;
    }

    if !invalid.is_empty() {
        write!(
            error_text,
            "\n\nInvalid account fields:\n{:?}",
            invalid.into_iter().format("\n")
        )?;
    }

    error_text.push_str("\n\nFix these problems and send the file again");

    Ok(Err(error_text))
}

#[inline]
fn user_from_vec(
    vector: Vec<u8>,
    existing_names: ahash::HashSet<String>,
) -> crate::Result<Result<User, String>> {
    let mut user: User = serde_json::from_slice(&vector)?;
    drop(vector);
    match process_accounts(&mut user.accounts, existing_names)? {
        Ok(()) => Ok(Ok(user)),
        Err(error_text) => Ok(Err(error_text)),
    }
}

#[inline]
async fn user_from_document(
    bot: &Throttle<Bot>,
    db: &Pool,
    document: Option<&Document>,
    user_id: u64,
) -> Result<User, Cow<'static, str>> {
    let (data, existing_names) = {
        let file = match validate_document(document) {
            Ok(document) => &document.file,
            Err(text) => return Err(Cow::Borrowed(text)),
        };

        let data =
            download_file(bot, file).map_err(|_| "Error getting existing account names. Try again");

        let existing_names = Account::get_names(user_id, db)
            .try_collect()
            .map_err(|_| "Error downloading the file. Try again");

        try_join!(data, existing_names)?
    };

    match spawn_blocking(|| user_from_vec(data, existing_names)).await {
        Ok(Ok(Ok(user))) => Ok(user),
        Ok(Ok(Err(error_text))) => Err(Cow::Owned(error_text)),
        _ => Err(Cow::Borrowed("Error parsing the json file. Try again")),
    }
}

/// Function to handle `GetUser` state. It doesn't actually validate anything
#[inline]
pub async fn get_user(
    bot: Throttle<Bot>,
    msg: Message,
    db: Pool,
    dialogue: MainDialogue,
    handler: PackagedHandler<User>,
) -> crate::Result<()> {
    let user_id = msg.from().ok_or(NoUserInfo)?.id.0;

    let mut handler = handler.lock().await;
    if handler.func.is_none() {
        let _ = dialogue.exit().await;
        return Err(HandlerUsed.into());
    }

    if msg.text().map(str::trim) == Some("/cancel") {
        dialogue.exit().await?;
        handler
            .previous
            .alter_message(&bot, "Successfully cancelled", deletion_markup(), None)
            .await?;
        return Ok(());
    }

    let user = match user_from_document(&bot, &db, msg.document(), user_id).await {
        Ok(user) => user,
        Err(error_text) => {
            handler
                .previous
                .alter_message(&bot, error_text, None, None)
                .await?;
            return Ok(());
        }
    };

    let previous = handler.previous;
    let func = handler.func.take().unwrap();
    drop(handler);

    if let Err(err) = func(bot, msg, db, dialogue.clone(), previous, user).await {
        let _ = dialogue.exit().await;
        return Err(err);
    }
    Ok(())
}