poc5.rs

use std::error::Error;
use std::str::FromStr;
use std::thread;
use std::time::Duration;

use bitcoin::hex::DisplayHex;
use bitcoin::{absolute, Amount, OutPoint, ScriptBuf, Sequence, Transaction, TxIn, TxOut};
use bitcoin::transaction::{Txid, Version};
use bitcoin::XOnlyPublicKey;

use bitcoincore_rpc::{Client, RpcApi};
use bitcoincore_rpc::json;

use clarity::types::Address as _;
use clarity::types::chainstate::StacksAddress;
use clarity::vm::types::{PrincipalData, StandardPrincipalData};
    
use emily_client::apis::configuration::{ApiKey, Configuration};
use emily_client::apis::deposit_api;
use emily_client::models;

use sbtc::deposits::{DepositScriptInputs, ReclaimScriptInputs};

use signer::config::Settings;
use signer::keys::SignerScriptPubKey;
use signer::storage::{postgres::PgStore, DbRead};

use secp256k1::PublicKey;

pub const DATABASE_URL: &str = "postgres://postgres:postgres@localhost:5432/signer";
pub const CONFIG_PATH: &str = "signer/src/config/default.toml";

pub const DEFAULT_STACKS_RECIPIENT: &str = "ST2SBXRBJJTH7GV5J93HJ62W2NRRQ46XYBK92Y039";

fn to_hex_string(input: &Vec<u8>) -> String {
    input.iter().map(|b| format!("{:02x}", b).to_string()).collect::<Vec<String>>().join("")
}

fn get_connection_pool(url: &str) -> sqlx::PgPool {
    sqlx::postgres::PgPoolOptions::new()
        .max_connections(1)
        .min_connections(1)
        .acquire_timeout(std::time::Duration::from_secs(5))
        .test_before_acquire(true)
        .connect_lazy(url)
        .unwrap()
}

async fn get_signers_script_pubkeys() -> Result<String, Box<dyn Error>> {
    let db = PgStore::connect(DATABASE_URL).await?;
    let signers_script_pubkeys = db.get_signers_script_pubkeys().await?;

    Ok(to_hex_string(&signers_script_pubkeys[0]))
}

async fn get_aggregate_key() -> Result<String, Box<dyn Error>> {
    let pool = get_connection_pool(DATABASE_URL);
    let aggregate_key = sqlx::query_scalar::<_, Vec<u8>>("SELECT aggregate_key FROM sbtc_signer.dkg_shares ORDER BY created_at DESC LIMIT 1;")
        .fetch_optional(&pool)
        .await?
        .unwrap();

    Ok(to_hex_string(&aggregate_key[1..].to_vec()))
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let settings = Settings::new(Some(CONFIG_PATH))?;

    let mut emily_api_endpoint = settings.emily.endpoints.first().unwrap().clone();
    let emily_api_key = if emily_api_endpoint.username().is_empty() {
        None
    } else {
        Some(ApiKey {
            prefix: None,
            key: emily_api_endpoint.username().to_string(),
        })
    };

    let _ = emily_api_endpoint.set_username("");
    let emily_client_config = Configuration {
        base_path: emily_api_endpoint.to_string().trim_end_matches("/").to_string(),
        api_key: emily_api_key,
        ..Default::default()
    };

    let bitcoin_url = format!("{}wallet/depositor", settings.bitcoin.rpc_endpoints.first().unwrap());
    let bitcoin_client = Client::new(
        &bitcoin_url,
        bitcoincore_rpc::Auth::UserPass("devnet".into(), "devnet".into()),
    )?;

    let aggregate_key = get_aggregate_key().await?;
    println!("aggregate_key: {:?}", aggregate_key);

    let signers_script_pubkeys = get_signers_script_pubkeys().await?;
    println!("signers_script_pubkeys: {:?}", signers_script_pubkeys);

    poc5(
        &bitcoin_client,
        &emily_client_config,
        aggregate_key,
    ).await?;

    Ok(())
}

async fn poc5(
    bitcoin_client: &Client,
    emily_config: &Configuration,
    aggregate_key: String,
) -> Result<(), Box<dyn Error>> {
    donation(bitcoin_client, aggregate_key.clone(), 40 * 100000000).await?;
    create_deposits_btc_and_emily(
        bitcoin_client,
        emily_config,
        aggregate_key,
        DEFAULT_STACKS_RECIPIENT.to_string(),
        50,
        10,
        42,
        20000,
        false,
    ).await?;

    println!("┏ Call Emliy to update all deposits to Confirmed");
    let get_deposits_resp = deposit_api::get_deposits(
        emily_config,
        models::Status::Pending,
        None,
        None,
    ).await?;

    let empty_fulfillment = models::Fulfillment::new(
        "".to_string(), 
        0,
        0,
        "".to_string(),
        0,
        "".to_string(),
    );

    let mut deposit_updates: Vec<models::DepositUpdate> = vec![];
    for pending_deposit in get_deposits_resp.deposits {
        deposit_updates.push(
            models::DepositUpdate {
                bitcoin_tx_output_index: pending_deposit.bitcoin_tx_output_index,
                bitcoin_txid: pending_deposit.bitcoin_txid,
                fulfillment: Some(Some(Box::new(empty_fulfillment.clone()))),
                last_update_block_hash: "".to_string(),
                last_update_height: 0,
                status: models::Status::Confirmed,
                status_message: "".to_string(),
            }
        );
    }

    deposit_api::update_deposits(
        emily_config,
        models::UpdateDepositsRequestBody::new(deposit_updates),
    ).await?;

    println!("┗ end");

    Ok(())
}

async fn wait_btc_tx_confirmed(
    bitcoin_client: &Client,
    txid: &Txid,
) -> Result<(), Box<dyn Error>> {
    println!("┃ wait for tx {:?} to be confirmed ...", txid);

    loop {
        println!("┃ ...");

        let tx = bitcoin_client.get_transaction(txid, None)?;
        if !tx.info.blockhash.is_none() {
            break;
        }

        thread::sleep(Duration::from_secs(10));
    }

    println!("┃ confirmed!");

    Ok(())
}

async fn create_deposits_btc_and_emily(
    bitcoin_client: &Client,
    emily_config: &Configuration,
    aggregate_key: String,
    stacks_recipient: String,
    lock_time: u32,
    number_of_deposits: u64,
    amount_per_deposit: u64,
    max_fee_per_deposit: u64,
    wait_tx_confirmed: bool,
) -> Result<(), Box<dyn Error>> {
    println!("┏ create_deposits_btc_and_emily");
    println!("┃ tx number of deposits: {:?}", number_of_deposits);

    let (unsigned_tx, deposit_script, reclaim_script) = create_bitcoin_deposit_transaction_multi_deposit(
        &bitcoin_client,
        aggregate_key,
        stacks_recipient,
        lock_time,
        number_of_deposits,
        amount_per_deposit,
        max_fee_per_deposit,
    )?;

    println!("┃ deposit script: {:?}", deposit_script.deposit_script().as_bytes().to_lower_hex_string());
    println!("┃ reclaim script: {:?}", reclaim_script.reclaim_script().as_bytes().to_lower_hex_string());

    let signed_tx = bitcoin_client.sign_raw_transaction_with_wallet(&unsigned_tx, None, None)?;
    // println!("Signed transaction: {:?}", hex::encode(&signed_tx.hex));
    let txid = bitcoin_client.send_raw_transaction(&signed_tx.hex)?;
    assert_eq!(txid, unsigned_tx.compute_txid());
    if wait_tx_confirmed {
        wait_btc_tx_confirmed(bitcoin_client, &txid).await?;
    }

    println!("┃ call emily to create all deposits");
    for i in 0..number_of_deposits {
        let _ = deposit_api::create_deposit(
            emily_config,
            models::CreateDepositRequestBody {
                bitcoin_tx_output_index: i as u32,
                bitcoin_txid: txid.to_string(),
                deposit_script: deposit_script.deposit_script().to_hex_string(),
                reclaim_script: reclaim_script.reclaim_script().to_hex_string(),
            },
        )
        .await?;
    }

    println!("┗ end");

    Ok(())
}

async fn donation(
    bitcoin_client: &Client,
    aggregate_key: String,
    donation_amount: u64,
) -> Result<(), Box<dyn Error>> {
    println!("┏ donation to aggregate_key");
    println!("┃ donation amount: {:?}", donation_amount);
    let opts = json::ListUnspentQueryOptions {
        minimum_amount: Some(Amount::from_sat(donation_amount)),
        ..Default::default()
    };
    let unspent = bitcoin_client
        .list_unspent(Some(6), None, None, None, Some(opts))?
        .into_iter()
        .next()
        .ok_or("list_unspent error")?;

    let pubkey = XOnlyPublicKey::from_str(&aggregate_key)
        .or_else(|_| PublicKey::from_str(&aggregate_key).map(XOnlyPublicKey::from))?;

    let relay_fee = 200;
    let unsigned_tx = Transaction {
        input: vec![TxIn {
            previous_output: OutPoint {
                txid: unspent.txid,
                vout: unspent.vout,
            },
            script_sig: Default::default(),
            sequence: Sequence::ZERO,
            witness: Default::default(),
        }],
        output: vec![
            TxOut {
                value: Amount::from_sat(donation_amount),
                script_pubkey: pubkey.signers_script_pubkey(),
            },
            TxOut {
                value: Amount::from_sat(unspent.amount.to_sat() - donation_amount - relay_fee),
                script_pubkey: unspent.script_pub_key, // Return the remaining BTC to unspent UTXO public key
            },
        ],
        version: Version::TWO,
        lock_time: absolute::LockTime::ZERO,
    };

    let signed_tx = bitcoin_client.sign_raw_transaction_with_wallet(&unsigned_tx, None, None)?;
    let txid = bitcoin_client.send_raw_transaction(&signed_tx.hex)?;
    wait_btc_tx_confirmed(bitcoin_client, &txid).await?;

    println!("┗ end");

    Ok(())
}

fn create_bitcoin_deposit_transaction_multi_deposit(
    bitcoin_client: &Client,
    aggregate_key: String,
    stacks_recipient: String,
    lock_time: u32,
    number_of_deposits: u64,
    amount_per_deposit: u64,
    max_fee_per_deposit: u64,
) -> Result<(Transaction, DepositScriptInputs, ReclaimScriptInputs), Box<dyn Error>> {
    let total_depsites_spend = (amount_per_deposit + max_fee_per_deposit) * number_of_deposits;
    let opts = json::ListUnspentQueryOptions {
        minimum_amount: Some(Amount::from_sat(total_depsites_spend)),
        ..Default::default()
    };
    let unspent = bitcoin_client
        .list_unspent(Some(6), None, None, None, Some(opts))?
        .into_iter()
        .next()
        .ok_or("list_unspent error")?;

    let pubkey = XOnlyPublicKey::from_str(&aggregate_key)
        .or_else(|_| PublicKey::from_str(&aggregate_key).map(XOnlyPublicKey::from))?;
    let deposit_script = DepositScriptInputs {
        signers_public_key: pubkey,
        max_fee: max_fee_per_deposit,
        recipient: PrincipalData::Standard(StandardPrincipalData::from(
            StacksAddress::from_string(&stacks_recipient).ok_or("StacksAddress::from_string error")?,
        )),
    };
    let reclaim_script = ReclaimScriptInputs::try_new(lock_time, ScriptBuf::new())?;

    let mut output: Vec<TxOut> = vec![];
    for _ in 0..number_of_deposits {
        output.push(TxOut {
            value: Amount::from_sat(amount_per_deposit + max_fee_per_deposit),
            script_pubkey: sbtc::deposits::to_script_pubkey(
                deposit_script.deposit_script(),
                reclaim_script.reclaim_script(),
            ),
        });
    }

    let relay_fee = 500+50*number_of_deposits;
    output.push(TxOut {
        value: Amount::from_sat(unspent.amount.to_sat() - total_depsites_spend - relay_fee),
        script_pubkey: unspent.script_pub_key, // Return the remaining BTC to unspent UTXO public key
    });

    let unsigned_tx = Transaction {
        input: vec![TxIn {
            previous_output: OutPoint {
                txid: unspent.txid,
                vout: unspent.vout,
            },
            script_sig: Default::default(),
            sequence: Sequence::ZERO,
            witness: Default::default(),
        }],
        output: output,
        version: Version::TWO,
        lock_time: absolute::LockTime::ZERO,
    };

    Ok((unsigned_tx, deposit_script, reclaim_script))
}