mod entity;

use entity::{account, prelude::Account};

use anyhow::Result;
use chacha20poly1305::{
    aead::{rand_core::RngCore, Aead, AeadCore, OsRng},
    ChaCha20Poly1305, KeyInit,
};
use dotenv::dotenv;
use pbkdf2::pbkdf2_hmac_array;
use sea_orm::{
    ActiveModelTrait, ActiveValue::Set, ColumnTrait, ConnectOptions, Database, EntityTrait,
    QueryFilter,
};
use sha2::Sha256;
use std::env;
use teloxide::{prelude::*, utils::command::BotCommands};

#[derive(BotCommands, Clone)]
#[command(rename_rule = "lowercase")]
enum Command {
    #[command()]
    Help,
    #[command(parse_with = "split")]
    AddAccount(String, String, String, String),
    #[command(parse_with = "split")]
    GetAccount(String, String),
}

use Command::*;

struct Cipher {
    chacha: ChaCha20Poly1305,
}

impl Cipher {
    fn new(password: &[u8], salt: &[u8]) -> Self {
        let key = pbkdf2_hmac_array::<Sha256, 32>(password, salt, 480000);

        Self {
            chacha: ChaCha20Poly1305::new(&key.into()),
        }
    }

    fn encrypt(&self, value: &[u8]) -> Result<Vec<u8>> {
        let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
        let mut result = self.chacha.encrypt(&nonce, value).unwrap();
        result.extend(nonce);
        Ok(result)
    }

    fn decrypt(&self, value: &[u8]) -> Result<Vec<u8>> {
        let (data, nonce) = value.split_at(value.len() - 12);
        assert!(nonce.len() == 12);
        self.chacha
            .decrypt(nonce.into(), data)
            .map_err(|err| err.into())
    }
}

async fn answer_command(db: ConnectOptions, bot: Bot, msg: Message, cmd: Command) -> Result<()> {
    let user_id = msg.from().unwrap().id.0;
    let db = Database::connect(db).await?;
    match cmd {
        Help => {
            bot.send_message(msg.chat.id, Command::descriptions().to_string())
                .await?;
        }
        AddAccount(name, login, password, master_pass) => {
            let mut salt = vec![0; 64];
            OsRng.fill_bytes(&mut salt);
            let cipher = Cipher::new(master_pass.as_ref(), &salt);
            let enc_login = Set(cipher.encrypt(login.as_ref())?);
            let enc_password = Set(cipher.encrypt(password.as_ref())?);
            let account = account::ActiveModel {
                name: Set(name),
                user_id: Set(user_id),
                salt: Set(salt),
                enc_login,
                enc_password,
            };
            account.insert(&db).await?;
            bot.send_message(msg.chat.id, "Success").await?;
        }
        GetAccount(name, master_pass) => {
            let account = Account::find()
                .filter(
                    account::Column::UserId
                        .eq(user_id)
                        .add(account::Column::Name.eq(&name)),
                )
                .one(&db)
                .await?
                .unwrap();
            let cipher = Cipher::new(master_pass.as_ref(), &account.salt);
            let login = String::from_utf8(cipher.decrypt(&account.enc_login)?)?;
            let password = String::from_utf8(cipher.decrypt(&account.enc_password)?)?;
            let message = format!("Account `{name}`\nLogin: `{login}`\nPassword: `{password}`");
            bot.send_message(msg.chat.id, message).await?;
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() {
    let _ = dotenv();
    pretty_env_logger::init();

    let token = env::var("TOKEN").unwrap();
    let database_url = env::var("DATABASE_URL").unwrap();
    let bot = Bot::new(token);
    let db = ConnectOptions::new(database_url);
    Dispatcher::builder(
        bot,
        Update::filter_message()
            .filter_command::<Command>()
            .endpoint(answer_command),
    )
    .dependencies(dptree::deps![db])
    .enable_ctrlc_handler()
    .build()
    .dispatch()
    .await;
}