bark-wallet 0.1.4

use anyhow::Context;
use bitcoin::{Amount, NetworkKind};
use bitcoin::hex::DisplayHex;
use log::{info, warn, error};

use ark::{VtxoPolicy, ProtocolEncoding};
use ark::arkoor::ArkoorDestination;
use ark::arkoor::package::{ArkoorPackageBuilder, ArkoorPackageCosignResponse};
use ark::vtxo::{Full, Vtxo, VtxoId};
use server_rpc::protos;

use crate::subsystem::Subsystem;
use crate::{ArkoorMovement, VtxoDelivery, MovementUpdate, Wallet, WalletVtxo};
use crate::movement::MovementDestination;
use crate::movement::manager::OnDropStatus;

/// The result of creating an arkoor transaction
pub struct ArkoorCreateResult {
	pub inputs: Vec<VtxoId>,
	pub created: Vec<Vtxo<Full>>,
	pub change: Vec<Vtxo<Full>>,
}

#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum ArkoorAddressError {
	#[error("Ark address is for different network")]
	NetworkMismatch,
	#[error("Ark address is for different server")]
	ServerMismatch,
	#[error("VTXO policy in address cannot be used for arkoor payment: {0:?}")]
	PolicyNotSupported(VtxoPolicy),
	#[error("No VTXO delivery mechanism provided in address")]
	NoDeliveryMechanism,
	#[error("Unknown delivery mechanism: {0}")]
	UnknownDeliveryMechanism(String),
	#[error("Other error: {0}")]
	Other(String),
}

impl Wallet {
	/// Validate if we can send arkoor payments to the given [ark::Address], for example an error
	/// will be returned if the given [ark::Address] belongs to a different server (see
	/// [ark::address::ArkId]).
	pub async fn validate_arkoor_address(&self, address: &ark::Address) -> Result<(), ArkoorAddressError> {
		let network = self.network().await
			.map_err(|e| ArkoorAddressError::Other(e.to_string()))?;
		let (_, ark_info) = self.require_server().await
			.map_err(|e| ArkoorAddressError::Other(e.to_string()))?;

		let network_kind = NetworkKind::from(network);
		if address.is_testnet() == network_kind.is_mainnet() {
			return Err(ArkoorAddressError::NetworkMismatch);
		}

		if !address.ark_id().is_for_server(ark_info.server_pubkey) {
			return Err(ArkoorAddressError::ServerMismatch);
		}

		// Not all policies are supported for sending arkoor
		match address.policy() {
			VtxoPolicy::Pubkey(_) => {},
			VtxoPolicy::ServerHtlcRecv(_) | VtxoPolicy::ServerHtlcSend(_) => {
				return Err(ArkoorAddressError::PolicyNotSupported(address.policy().clone()));
			}
		}

		if address.delivery().is_empty() {
			return Err(ArkoorAddressError::NoDeliveryMechanism);
		}
		// We first see if we know any of the deliveries, if not, we will log
		// the unknown onces.
		// We do this in two parts because we shouldn't log unknown ones if there is one known.
		if !address.delivery().iter().any(|d| !d.is_unknown()) {
			for d in address.delivery() {
				if let VtxoDelivery::Unknown { delivery_type, data } = d {
					info!("Unknown delivery in address: type={:#x}, data={}",
						delivery_type, data.as_hex(),
					);
				}
			}
		}

		Ok(())
	}

	pub(crate) async fn create_checkpointed_arkoor_with_vtxos(
		&self,
		arkoor_dest: ArkoorDestination,
		inputs: impl IntoIterator<Item = WalletVtxo>,
	) -> anyhow::Result<ArkoorCreateResult> {
		// Find vtxos to cover
		let (mut srv, _) = self.require_server().await?;
		let (input_ids, inputs) = inputs.into_iter()
			.map(|v| (v.id(), v))
			.collect::<(Vec<_>, Vec<_>)>();

		// Pre-register the input chains so the post-cosign register call
		// for the outputs finds a signed chain anchor:
		// register_vtxo_transactions validates a vtxo against its anchor's
		// signed_tx in the DB, and boarded inputs sit unsigned in
		// virtual_transaction (see register_board) until a
		// register_vtxo_transactions call backfills them.
		self.register_vtxo_transactions_with_server(&inputs).await
			.context("failed to register arkoor input vtxo transactions with server")?;

		// Peek at a potential change keypair without storing it yet.
		// We'll only store it if change is actually created.
		let (change_keypair, change_key_index) = self.peek_next_keypair().await?;
		let change_pubkey = change_keypair.public_key();

		if arkoor_dest.policy.user_pubkey() == change_pubkey {
			bail!("Cannot create arkoor to same address as change");
		}

		let mut user_keypairs = vec![];
		for vtxo in &inputs {
			user_keypairs.push(self.get_vtxo_key(vtxo).await?);
		}

		let builder = ArkoorPackageBuilder::new_single_output_with_checkpoints(
			inputs.into_iter().map(|v| v.vtxo),
			arkoor_dest.clone(),
			VtxoPolicy::new_pubkey(change_pubkey),
		)
			.context("Failed to construct arkoor package")?
			.generate_user_nonces(&user_keypairs)
			.context("invalid nb of keypairs")?;

		let cosign_request = protos::ArkoorPackageCosignRequest::from(
			builder.cosign_request(),
		);

		let response = srv.client.request_arkoor_cosign(cosign_request).await
			.context("server failed to cosign arkoor")?
			.into_inner();

		let cosign_responses = ArkoorPackageCosignResponse::try_from(response)
			.context("Failed to parse cosign response from server")?;

		let vtxos = builder
			.user_cosign(&user_keypairs, cosign_responses)
			.context("Failed to cosign vtxos")?
			.build_signed_vtxos();

		// divide between change and destination
		let (dest, change) = vtxos.into_iter()
			.partition::<Vec<_>, _>(|v| *v.policy() == arkoor_dest.policy);

		if !change.is_empty() {
			// Change was created, so now store the keypair
			self.db.store_vtxo_key(change_key_index, change_pubkey).await?;
		}

		Ok(ArkoorCreateResult {
			inputs: input_ids,
			created: dest,
			change,
		})
	}

	/// Makes an out-of-round payment to the given [ark::Address]. This does not require waiting for
	/// a round, so it should be relatively instantaneous.
	///
	/// If the [Wallet] doesn't contain a VTXO larger than the given [Amount], multiple payments
	/// will be chained together, resulting in the recipient receiving multiple VTXOs.
	///
	/// Note that a change [Vtxo] may be created as a result of this call. With each payment these
	/// will become more uneconomical to unilaterally exit, so you should eventually refresh them
	/// with [Wallet::refresh_vtxos] or periodically call [Wallet::maintenance_refresh].
	pub async fn send_arkoor_payment(
		&self,
		destination: &ark::Address,
		amount: Amount,
	) -> anyhow::Result<Vec<Vtxo<Full>>> {
		let (mut srv, _) = self.require_server().await?;

		self.validate_arkoor_address(&destination).await
			.context("address validation failed")?;

		let negative_amount = -amount.to_signed().context("Amount out-of-range")?;

		let dest = ArkoorDestination { total_amount: amount, policy: destination.policy().clone() };
		let inputs = self.select_vtxos_to_cover(dest.total_amount).await?;
		let arkoor = self.create_checkpointed_arkoor_with_vtxos(dest.clone(), inputs.into_iter())
			.await.context("failed to create arkoor transaction")?;

		// Register destination and change after cosign so the server
		// gets signed_tx rows for the shared checkpoint and both leaves:
		// the watchman needs the checkpoint signed to rebroadcast on
		// exit, and a later guarded spend (offboard, lightning pay) of
		// either leaf needs the full chain signed. Failure is logged
		// and swallowed: the receiver also re-registers on receive, and
		// later spends will retry registration if the server still
		// needs it.
		let to_register = arkoor.created.iter()
			.chain(arkoor.change.iter())
			.collect::<Vec<_>>();
		if let Err(e) = self.register_vtxo_transactions_with_server(&to_register).await {
			warn!("Failed to register arkoor output vtxo transactions with server: {:#}", e);
		}

		let mut movement = self.movements.new_guarded_movement_with_update(
			Subsystem::ARKOOR,
			ArkoorMovement::Send.to_string(),
			OnDropStatus::Failed,
			MovementUpdate::new()
				.intended_and_effective_balance(negative_amount)
				.consumed_vtxos(&arkoor.inputs)
				.sent_to([MovementDestination::ark(destination.clone(), amount)])
		).await?;

		let mut delivered = false;
		for delivery in destination.delivery() {
			match delivery {
				VtxoDelivery::ServerMailbox { blinded_id } => {
					let req = protos::mailbox_server::PostArkoorMessageRequest {
						blinded_id: blinded_id.to_vec(),
						vtxos: arkoor.created.iter().map(|v| v.serialize().to_vec()).collect(),
					};

					if let Err(e) = srv.mailbox_client.post_arkoor_message(req).await {
						error!("Failed to post the vtxos to the destination's mailbox: '{:#}'", e);
						//NB we will continue to at least not lose our own change
					} else {
						delivered = true;
					}
				},
				VtxoDelivery::Unknown { delivery_type, data } => {
					error!("Unknown delivery type {} for arkoor payment: {}", delivery_type, data.as_hex());
				},
				_ => {
					error!("Unsupported delivery type for arkoor payment: {:?}", delivery);
				}
			}
		}
		self.mark_vtxos_as_spent(&arkoor.inputs).await?;
		if !arkoor.change.is_empty() {
			self.store_spendable_vtxos(&arkoor.change).await?;
			movement.apply_update(MovementUpdate::new().produced_vtxos(arkoor.change)).await?;
		}

		if delivered {
			movement.success().await?;
		} else {
			bail!("Failed to deliver arkoor vtxos to any destination");
		}

		Ok(arkoor.created)
	}
}