brain/src/db/general.rs

use super::Error;
use crate::constants::{ACCOUNT, BAN, KICK, MIN_ESCROW, TOKEN_DECIMAL, VM_NODE};
use crate::db::prelude::*;
use crate::old_brain;
use serde::{Deserialize, Serialize};
use surrealdb::engine::remote::ws::Client;
use surrealdb::sql::Datetime;
use surrealdb::{RecordId, Surreal};

#[derive(Debug, Serialize, Deserialize)]
pub struct Account {
    pub id: RecordId,
    pub balance: u64,
    pub tmp_locked: u64,
    pub escrow: u64,
    pub email: String,
}

impl Account {
    pub async fn get(db: &Surreal<Client>, address: &str) -> Result<Self, Error> {
        let id = (ACCOUNT, address);
        let account: Option<Self> = db.select(id).await?;
        let account = match account {
            Some(account) => account,
            None => {
                Self { id: id.into(), balance: 0, tmp_locked: 0, escrow: 0, email: String::new() }
            }
        };
        Ok(account)
    }

    pub async fn get_or_create(db: &Surreal<Client>, address: &str) -> Result<Self, Error> {
        let id = (ACCOUNT, address);

        match db.select(id).await? {
            Some(account) => Ok(account),
            None => {
                let account: Option<Self> = db.create(id).await?;
                account.ok_or(Error::FailedToCreateDBEntry(ACCOUNT.to_string()))
            }
        }
    }

    pub async fn airdrop(db: &Surreal<Client>, account: &str, tokens: u64) -> Result<(), Error> {
        let tokens = tokens.saturating_mul(1_000_000_000);
        let _ = db
            .query(format!("upsert account:{account} SET balance = (balance || 0) + {tokens};"))
            .await?;
        Ok(())
    }

    pub async fn save(self, db: &Surreal<Client>) -> Result<Option<Self>, Error> {
        let account: Option<Self> = db.upsert(self.id.clone()).content(self).await?;
        Ok(account)
    }

    pub async fn operator_reg(
        db: &Surreal<Client>,
        wallet: &str,
        email: &str,
        escrow: u64,
    ) -> Result<(), Error> {
        if escrow < MIN_ESCROW {
            return Err(Error::MinimalEscrow);
        }
        let mut op_account = Self::get(db, wallet).await?;
        let escrow = escrow.saturating_mul(TOKEN_DECIMAL);
        let op_total_balance = op_account.balance.saturating_add(op_account.escrow);
        if op_total_balance < escrow {
            return Err(Error::InsufficientFunds);
        }
        op_account.email = email.to_string();
        op_account.balance = op_total_balance.saturating_sub(escrow);
        op_account.escrow = escrow;
        op_account.save(db).await?;
        Ok(())
    }

    pub async fn slash_account(
        db: &Surreal<Client>,
        account: &str,
        slash_amount: u64,
    ) -> Result<(), Error> {
        let tx_query = "
        BEGIN TRANSACTION;
            LET $account = $account_input;
        
            UPDATE $account SET escrow -= $slash_amount;
            IF $account.escrow < 0 {{
                THROW 'Insufficient escrow.'
            }};

        COMMIT TRANSACTION;";

        let mut query_resp = db
            .query(tx_query)
            .bind(("account_input", RecordId::from((ACCOUNT, account))))
            .bind(("slash_amount", slash_amount.saturating_mul(TOKEN_DECIMAL)))
            .await?;

        log::trace!("query_resp: {query_resp:?}");

        let query_error = query_resp.take_errors();
        if !query_error.is_empty() {
            log::error!("slash_account query error: {query_error:?}");
            return Err(Error::FailedToSlashOperator(account.to_string()));
        }
        Ok(())
    }

    pub async fn list_accounts(db: &Surreal<Client>) -> Result<Vec<Self>, Error> {
        let accounts: Vec<Account> = db.select(ACCOUNT).await?;
        Ok(accounts)
    }
}

impl Account {
    pub async fn is_banned_by_node(
        db: &Surreal<Client>,
        user: &str,
        node: &str,
    ) -> Result<bool, Error> {
        let mut query_response = db
            .query(format!(
                "(select operator->ban[0] as ban
                    from $vm_node_input
                    where operator->ban->account contains $user_account_input
                ).ban;"
            ))
            .bind(("vm_node_input", RecordId::from((VM_NODE, node))))
            .bind(("user_account_input", RecordId::from((ACCOUNT, user))))
            .query(format!(
                "(select operator->ban[0] as ban
                    from $app_node_input
                    where operator->ban->account contains $user_account_input
                ).ban;"
            ))
            .bind(("app_node_input", RecordId::from((APP_NODE, node))))
            .bind(("user_account_input", RecordId::from((ACCOUNT, user))))
            .await?;

        let vm_node_ban: Option<Ban> = query_response.take(0).unwrap();
        let app_node_ban: Option<Ban> = query_response.take(1).unwrap();

        log::trace!("vm_node_ban: {vm_node_ban:?}");
        log::trace!("app_node_ban: {app_node_ban:?}");

        Ok(vm_node_ban.is_some() || app_node_ban.is_some())
    }
}

impl From<&old_brain::BrainData> for Vec<Account> {
    fn from(old_data: &old_brain::BrainData) -> Self {
        let mut accounts = Vec::new();
        for old_account in old_data.accounts.iter() {
            let mut a = Account {
                id: RecordId::from((ACCOUNT, old_account.key())),
                balance: old_account.value().balance,
                tmp_locked: old_account.value().tmp_locked,
                escrow: 0,
                email: String::new(),
            };
            if let Some(operator) = old_data.operators.get(old_account.key()) {
                a.escrow = operator.escrow;
                a.email = operator.email.clone();
            }
            accounts.push(a);
        }
        accounts
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Ban {
    #[serde(rename = "in")]
    pub from_account: RecordId,
    #[serde(rename = "out")]
    pub to_account: RecordId,
    pub created_at: Datetime,
}

impl Ban {
    pub async fn get_record(
        db: &Surreal<Client>,
        op_wallet: &str,
        user_wallet: &str,
    ) -> Result<Option<Self>, Error> {
        let query =
            format!("SELECT * FROM {BAN} WHERE in = $operator_input AND out = $user_input;");
        let mut response = db
            .query(query)
            .bind(("operator_input", RecordId::from((ACCOUNT, op_wallet))))
            .bind(("user_input", RecordId::from((ACCOUNT, user_wallet))))
            .await?;
        let ban_record: Option<Self> = response.take(0)?;
        Ok(ban_record)
    }

    pub async fn create(
        db: &Surreal<Client>,
        op_wallet: &str,
        user_wallet: &str,
    ) -> Result<(), Error> {
        if Self::get_record(db, op_wallet, user_wallet).await?.is_some() {
            log::error!("User {user_wallet} is already banned by {op_wallet}");
            return Err(Error::AlreadyBanned(op_wallet.to_string()));
        }
        let _: Vec<Self> = db
            .insert(BAN)
            .relation(Self {
                from_account: RecordId::from((ACCOUNT, op_wallet)),
                to_account: RecordId::from((ACCOUNT, user_wallet)),
                created_at: Default::default(),
            })
            .await?;
        Ok(())
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Kick {
    id: RecordId,
    #[serde(rename = "in")]
    from_account: RecordId,
    #[serde(rename = "out")]
    to_account: RecordId,
    created_at: Datetime,
    reason: String,
    contract: RecordId,
    node: RecordId,
}

impl Kick {
    pub async fn kicked_in_a_day(db: &Surreal<Client>, account: &str) -> Result<Vec<Self>, Error> {
        let mut result = db
            .query(format!(
                "select * from {KICK} where out = {ACCOUNT}:{account} and created_at > time::now() - 24h;"
            ))
            .await?;
        let kicks: Vec<Self> = result.take(0)?;
        Ok(kicks)
    }

    pub async fn submit(self, db: &Surreal<Client>) -> Result<(), Error> {
        let _: Vec<Self> = db.insert(KICK).relation(self).await?;
        Ok(())
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Report {
    #[serde(rename = "in")]
    from_account: RecordId,
    #[serde(rename = "out")]
    to_node: RecordId,
    created_at: Datetime,
    pub reason: String,
    pub contract_id: String,
}

impl Report {
    // TODO: test this functionality and remove this comment
    pub async fn create(
        db: &Surreal<Client>,
        from_account: RecordId,
        to_node: RecordId,
        reason: String,
        contract_id: String,
    ) -> Result<(), Error> {
        let _: Vec<Self> = db
            .insert("report")
            .relation(Report {
                from_account,
                to_node,
                created_at: Datetime::default(),
                reason,
                contract_id,
            })
            .await?;
        Ok(())
    }
}

/// This is the operator obtained from the DB,
/// however the relation is defined using OperatorRelation
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Operator {
    pub account: RecordId,
    pub app_nodes: u64,
    pub vm_nodes: u64,
    pub email: String,
    pub escrow: u64,
    pub reports: u64,
}

impl Operator {
    pub async fn list(db: &Surreal<Client>) -> Result<Vec<Self>, Error> {
        let mut result = db
            .query(
                "array::distinct(array::flatten( [
                    (select operator from vm_node group by operator).operator,
                    (select operator from app_node group by operator).operator
                ]));"
                    .to_string(),
            )
            .await?;
        let operator_accounts: Vec<RecordId> = result.take(0)?;
        let mut operators: Vec<Self> = Vec::new();
        for account in operator_accounts.iter() {
            if let Some(operator) = Self::inspect(db, &account.key().to_string()).await? {
                operators.push(operator);
            }
        }
        Ok(operators)
    }

    pub async fn inspect(db: &Surreal<Client>, account: &str) -> Result<Option<Self>, Error> {
        let mut result = db
            .query(format!(
                "LET $vm_nodes = (select id from vm_node where operator = account:{account}).id;
                LET $app_nodes = (select id from app_node where operator = account:{account}).id;
                select *,
                    id as account,
                    email,
                    escrow,
                    $vm_nodes.len() as vm_nodes,
                    $app_nodes.len() as app_nodes,
                    (select id from report where $vm_nodes contains out).len() +
                    (select id from report where $app_nodes contains out).len()
                    as reports
                from account where id = account:{account};"
            ))
            .await?;
        let operator: Option<Self> = result.take(2)?;
        Ok(operator)
    }

    pub async fn inspect_nodes(
        db: &Surreal<Client>,
        account: &str,
    ) -> Result<(Option<Self>, Vec<VmNodeWithReports>, Vec<AppNodeWithReports>), Error> {
        let operator = Self::inspect(db, account).await?;
        let mut result = db
            .query(format!(
                "select *, operator, <-report.* as reports from vm_node
                 where operator = account:{account};"
            ))
            .query(format!(
                "select *, operator, <-report.* as reports from app_node
                 where operator = account:{account};"
            ))
            .await?;
        let vm_nodes: Vec<VmNodeWithReports> = result.take(0)?;
        let app_nodes: Vec<AppNodeWithReports> = result.take(1)?;

        Ok((operator, vm_nodes, app_nodes))
    }
}

pub async fn kick_contract(
    db: &Surreal<Client>,
    operator_wallet: &str,
    contract_uuid: &str,
    reason: &str,
) -> Result<u64, Error> {
    let (contract_id, operator_id, admin_id, app_or_vm) =
        if let Some(active_vm) = ActiveVmWithNode::get_by_uuid(db, contract_uuid).await? {
            (active_vm.id, active_vm.vm_node.operator, active_vm.admin, "vm")
        } else if let Some(active_app) = ActiveAppWithNode::get_by_uuid(db, contract_uuid).await? {
            (active_app.id, active_app.app_node.operator, active_app.admin, "app")
        } else {
            return Err(Error::ContractNotFound);
        };

    if operator_id.key().to_string() != operator_wallet {
        return Err(Error::AccessDenied);
    }

    log::debug!("Kicking contract {contract_id} by operator {operator_id} for reason: '{reason}'",);

    let transaction_query = format!(
        "
        BEGIN TRANSACTION;
            LET $contract = {contract_id};
            LET $operator_account = {operator_id};
            LET $reason = $reason_str_input;
            LET $contract_id = record::id($contract.id);
            LET $admin = $contract.in;
            LET $node = $contract.out;

            LET $active_contract = (select * from $contract)[0];
            LET $deleted_contract = $active_contract.patch([{{
                'op': 'replace',
                'path': 'id',
                'value': type::record('deleted_{app_or_vm}:' + $contract_id)
            }}]);
            LET $deleted_contract = (INSERT RELATION INTO deleted_{app_or_vm} ( $deleted_contract ) RETURN AFTER)[0];

            -- calculating refund minutes
            LET $one_week_minutes = duration::mins(1w);
            LET $uncollected_minutes =  (time::now() - $active_contract.collected_at).mins();
            
            LET $minutes_to_refund = if $uncollected_minutes > $one_week_minutes {{
                $one_week_minutes;
            }} ELSE {{
                $uncollected_minutes;
            }};

            -- calculating refund amount
            LET $prince_per_minute = fn::{app_or_vm}_price_per_minute($active_contract.id);
            
            LET $refund_amount = $prince_per_minute * $minutes_to_refund;

            LET $refund =  IF SELECT * FROM {KICK} WHERE out = $admin.id AND created_at > time::now() - 24h {{
                0
            }} ELSE IF $operator_account.escrow <= $refund_amount {{
                $operator_account.escrow
            }} ELSE {{
                $refund_amount;
            }};

            RELATE $operator_account->{KICK}->$admin
            SET id = $contract_id,
                reason = $reason,
                contract = $deleted_contract.id,
                node = $node.id,
                created_at = time::now()
            ;
            DELETE $active_contract.id;

            -- update balances
            UPDATE $operator_account SET escrow -= $refund;
            IF $operator_account.escrow < 0 {{
                THROW 'Insufficient funds.'
            }};
            UPDATE $admin SET balance += $refund;

            $refund;

        COMMIT TRANSACTION;
        ",
    );

    log::trace!("kick_contract transaction_query: {}", &transaction_query);

    let mut query_res =
        db.query(transaction_query).bind(("reason_str_input", reason.to_string())).await?;

    log::trace!("transaction_query response: {:?}", &query_res);

    let query_error = query_res.take_errors();
    if !query_error.is_empty() {
        log::error!("kick_contract query error: {query_error:?}");
        return Err(Error::FailedKickContract(contract_id.to_string()));
    }

    let refunded: Option<u64> = query_res.take(20)?;
    let refunded_amount = refunded.ok_or(Error::FailedToCreateDBEntry("Refund".to_string()))?;
    log::info!("Refunded: {refunded_amount} to {admin_id}");

    Ok(refunded_amount)
}