bark-wallet 0.1.4

use anyhow::Context;
use bitcoin::{Amount, SignedAmount, Transaction, Txid};
use bitcoin::hashes::Hash;
use bitcoin::hex::DisplayHex;
use bitcoin::secp256k1::Keypair;
use log::info;

use ark::{musig, ProtocolEncoding, Vtxo, VtxoPolicy};
use ark::arkoor::ArkoorDestination;
use ark::attestations::OffboardRequestAttestation;
use ark::fees::{validate_and_subtract_fee_min_dust, VtxoFeeInfo};
use ark::offboard::{OffboardForfeitContext, OffboardRequest};
use ark::vtxo::{Full, VtxoRef};
use bitcoin_ext::P2TR_DUST;
use server_rpc::{protos, ServerConnection, TryFromBytes};

use crate::movement::manager::OnDropStatus;
use crate::vtxo::VtxoState;
use crate::{Wallet, WalletVtxo};
use crate::movement::update::MovementUpdate;
use crate::movement::{MovementDestination, MovementStatus};
use crate::persist::models::PendingOffboard;
use crate::subsystem::{OffboardMovement, Subsystem};


impl Wallet {
	async fn offboard_inner(
		&self,
		srv: &mut ServerConnection,
		vtxos: &[impl AsRef<Vtxo<Full>>],
		vtxo_keys: &[Keypair],
		req: &OffboardRequest,
	) -> anyhow::Result<Transaction> {
		// Register VTXO transaction chains with server before offboarding
		self.register_vtxo_transactions_with_server(&vtxos).await?;

		let input_ids = vtxos.iter().map(|v| v.as_ref().id()).collect::<Vec<_>>();
		let prep_resp = srv.client.prepare_offboard(protos::PrepareOffboardRequest {
			offboard: Some(req.into()),
			input_vtxo_ids: input_ids.iter()
				.map(|id| id.to_bytes().to_vec())
				.collect(),
			attestation: vtxo_keys.iter()
				.map(|k| OffboardRequestAttestation::new(req, &input_ids, k).serialize())
				.collect(),
		}).await.context("prepare offboard request failed")?.into_inner();
		let unsigned_offboard_tx = bitcoin::consensus::deserialize::<Transaction>(
			&prep_resp.offboard_tx,
		).with_context(|| format!(
			"received invalid unsigned offboard tx from server: {}", prep_resp.offboard_tx.as_hex(),
		))?;
		let offboard_txid = unsigned_offboard_tx.compute_txid();
		info!("Received unsigned offboard tx {} from server", offboard_txid);
		let forfeit_cosign_nonces = prep_resp.forfeit_cosign_nonces.into_iter().map(|n| {
			Ok(musig::PublicNonce::from_bytes(&n)
				.context("received invalid public cosign nonce from server")?)
		}).collect::<anyhow::Result<Vec<_>>>()?;

		let ctx = OffboardForfeitContext::new(&vtxos, &unsigned_offboard_tx);
		ctx.validate_offboard_tx(&req).context("received invalid offboard tx from server")?;

		let sigs = ctx.user_sign_forfeits(&vtxo_keys, &forfeit_cosign_nonces);

		let finish_resp = srv.client.finish_offboard(protos::FinishOffboardRequest {
			offboard_txid: offboard_txid.as_byte_array().to_vec(),
			user_nonces: sigs.public_nonces.iter().map(|n| n.serialize().to_vec()).collect(),
			partial_signatures: sigs.partial_signatures.iter()
				.map(|s| s.serialize().to_vec())
				.collect(),
		}).await.context("error sending offboard forfeit signatures to server")?.into_inner();

		let signed_offboard_tx = bitcoin::consensus::deserialize::<Transaction>(
			&finish_resp.signed_offboard_tx,
		).with_context(|| format!(
			"received invalid offboard tx from server: {}", finish_resp.signed_offboard_tx.as_hex(),
		))?;
		if signed_offboard_tx.compute_txid() != offboard_txid {
			bail!("Signed offboard tx received from server is different from \
				unsigned tx we forfeited for: unsigned={}, signed={}",
				prep_resp.offboard_tx.as_hex(), finish_resp.signed_offboard_tx.as_hex(),
			);
		}

		// we don't accept the tx if our mempool doesn't accept it, it might be a double spend
		self.chain.broadcast_tx(&signed_offboard_tx).await.with_context(|| format!(
			"error broadcasting offboard tx {} (tx={})",
			offboard_txid, finish_resp.signed_offboard_tx.as_hex(),
		))?;

		Ok(signed_offboard_tx)
	}

	/// Send to an onchain address using your offchain balance
	pub async fn send_onchain(
		&self,
		destination: bitcoin::Address,
		amount: Amount,
	) -> anyhow::Result<Txid> {
		if amount < P2TR_DUST {
			bail!("it doesn't make sense to send dust");
		}

		let (mut srv, ark) = self.require_server().await?;

		let destination_spk = destination.script_pubkey();
		let (vtxos, fee) = self.select_vtxos_to_cover_with_fee(amount, |a, v| {
			ark.fees.offboard.calculate(&destination_spk, a, ark.offboard_feerate, v)
				.ok_or_else(|| anyhow!("failed to calculate offboard fee for {}", a))
		}).await?;
		let required_amount = amount + fee;

		info!("We can only offboard whole VTXOs, so we will make an arkoor tx first...");

		// this will be the key that holds the temporary vtxos we will offboard
		let offboard_pubkey = self.derive_store_next_keypair().await
			.context("failed to create new keypair")?.0;
		let offboard_dest = ArkoorDestination {
			total_amount: required_amount,
			policy: VtxoPolicy::new_pubkey(offboard_pubkey.public_key()),
		};
		let arkoor = self.create_checkpointed_arkoor_with_vtxos(offboard_dest, vtxos.into_iter())
			.await.context("error trying to prepare offboard VTXOs with an arkoor tx")?;

		self.store_spendable_vtxos(&arkoor.change).await
			.context("error storing change VTXOs from preparatory arkoor")?;
		self.store_locked_vtxos(&arkoor.created, None).await
			.context("error storing new VTXOs (locked) from preparatory arkoor")?;
		self.mark_vtxos_as_spent(&arkoor.inputs).await
			.context("error marking used input VTXOs as spent")?;

		let mut movement = self.movements.new_guarded_movement_with_update(
			Subsystem::OFFBOARD,
			OffboardMovement::SendOnchain.to_string(),
			OnDropStatus::Failed,
			MovementUpdate::new()
				.intended_balance(-amount.to_signed()?)
				.effective_balance(-required_amount.to_signed()?)
				.fee(fee)
				.consumed_vtxos(&arkoor.inputs)
				.produced_vtxos(&arkoor.change)
				.metadata([(
					"offboard_vtxos".into(),
					serde_json::to_value(
						arkoor.created.iter().map(|v| v.id()).collect::<Vec<_>>(),
					).expect("offboard_vtxos can serde"),
				)])
				.sent_to([MovementDestination::bitcoin(destination.clone(), amount)])
		).await?;
		let state = VtxoState::Locked { movement_id: Some(movement.id()) };
		self.set_vtxo_states(&arkoor.created, &state, &[]).await
			.context("error setting movement id on locked VTXOs")?;

		// now perform the offboard
		let vtxos = arkoor.created;

		let req = OffboardRequest {
			script_pubkey: destination_spk.clone(),
			net_amount: amount,
			deduct_fees_from_gross_amount: false,
			fee_rate: ark.offboard_feerate,
		};
		let vtxo_keys = vec![offboard_pubkey; vtxos.len()];

		let signed_offboard_tx = self.offboard_inner(&mut srv, &vtxos, &vtxo_keys, &req).await
			.context("error performing offboard")?;

		movement.apply_update(MovementUpdate::new()
			.metadata(OffboardMovement::metadata(&signed_offboard_tx))
		).await.context("error updating movement")?;

		if self.config.offboard_required_confirmations == 0 {
			// No confirmation required — mark VTXOs as spent and succeed immediately
			for vtxo in &vtxos {
				self.db.update_vtxo_state_checked(
					vtxo.id(),
					VtxoState::Spent,
					&[crate::vtxo::VtxoStateKind::Locked],
				).await.context("error marking vtxo as spent")?;
			}
			movement.success().await
				.context("error finishing movement")?;
		} else {
			// Store as pending offboard — don't mark success until confirmed on chain
			let vtxo_ids = vtxos.iter().map(|v| v.id()).collect::<Vec<_>>();
			self.db.store_pending_offboard(&PendingOffboard {
				movement_id: movement.id(),
				offboard_txid: signed_offboard_tx.compute_txid(),
				offboard_tx: signed_offboard_tx.clone(),
				vtxo_ids,
				destination: destination.to_string(),
				created_at: chrono::Local::now(),
			}).await.context("error storing pending offboard")?;

			// Disarm the guard so it doesn't auto-fail the movement on drop
			movement.stop();
		}

		Ok(signed_offboard_tx.compute_txid())
	}

	async fn offboard(
		&self,
		vtxos: Vec<WalletVtxo>,
		destination: bitcoin::Address,
	) -> anyhow::Result<Txid> {
		let (mut srv, ark) = self.require_server().await?;
		let tip = self.chain.tip().await?;

		let destination_spk = destination.script_pubkey();
		let vtxos_amount = vtxos.iter().map(|v| v.amount()).sum::<Amount>();
		let fee = ark.fees.offboard.calculate(
			&destination_spk,
			vtxos_amount,
			ark.offboard_feerate,
			vtxos.iter().map(|v| VtxoFeeInfo::from_vtxo_and_tip(v, tip)),
		).context("error calculating offboard fee")?;
		let net_amount = validate_and_subtract_fee_min_dust(vtxos_amount, fee)?;
		let vtxo_keys = {
			let mut keys = Vec::with_capacity(vtxos.len());
			for v in &vtxos {
				keys.push(self.get_vtxo_key(v).await?);
			}
			keys
		};

		let req = OffboardRequest {
			script_pubkey: destination_spk.clone(),
			net_amount,
			deduct_fees_from_gross_amount: true,
			fee_rate: ark.offboard_feerate,
		};

		let signed_offboard_tx = self.offboard_inner(&mut srv, &vtxos, &vtxo_keys, &req).await
			.context("error performing offboard")?;

		// Lock VTXOs instead of marking them as spent
		let vtxo_ids = vtxos.iter().map(|v| v.vtxo_id()).collect::<Vec<_>>();
		let effective_amt = -SignedAmount::try_from(vtxos_amount)
			.expect("can't have this many vtxo sats");
		let destination_str = destination.to_string();
		let movement_id = self.movements.new_movement_with_update(
			Subsystem::OFFBOARD,
			OffboardMovement::Offboard.to_string(),
			MovementUpdate::new()
				.intended_balance(effective_amt)
				.effective_balance(effective_amt)
				.fee(fee)
				.consumed_vtxos(&vtxos)
				.sent_to([MovementDestination::bitcoin(destination, net_amount)])
				.metadata(OffboardMovement::metadata(&signed_offboard_tx)),
		).await?;

		self.lock_vtxos(&vtxos, Some(movement_id)).await?;

		if self.config.offboard_required_confirmations == 0 {
			// No confirmation required — mark VTXOs as spent and succeed immediately
			for vtxo in &vtxos {
				self.db.update_vtxo_state_checked(
					vtxo.vtxo_id(),
					VtxoState::Spent,
					&[crate::vtxo::VtxoStateKind::Locked],
				).await.context("error marking vtxo as spent")?;
			}
			self.movements.finish_movement(
				movement_id,
				MovementStatus::Successful,
			).await.context("error finishing movement")?;
		} else {
			// Store as pending offboard — wait for on-chain confirmation
			self.db.store_pending_offboard(&PendingOffboard {
				movement_id,
				offboard_txid: signed_offboard_tx.compute_txid(),
				offboard_tx: signed_offboard_tx.clone(),
				vtxo_ids,
				destination: destination_str,
				created_at: chrono::Local::now(),
			}).await.context("error storing pending offboard")?
		}

		Ok(signed_offboard_tx.compute_txid())
	}

	/// Offboard all VTXOs to a given [bitcoin::Address].
	pub async fn offboard_all(&self, address: bitcoin::Address) -> anyhow::Result<Txid> {
		let input_vtxos = self.spendable_vtxos().await?;
		Ok(self.offboard(input_vtxos, address).await?)
	}

	/// Offboard the given VTXOs to a given [bitcoin::Address].
	pub async fn offboard_vtxos<V: VtxoRef>(
		&self,
		vtxos: impl IntoIterator<Item = V>,
		address: bitcoin::Address,
	) -> anyhow::Result<Txid> {
		let mut input_vtxos = vec![];
		for v in vtxos {
			let id = v.vtxo_id();
			let vtxo = match self.db.get_wallet_vtxo(id).await? {
				Some(vtxo) => vtxo,
				_ => bail!("cannot find requested vtxo: {}", id),
			};
			input_vtxos.push(vtxo);
		}

		Ok(self.offboard(input_vtxos, address).await?)
	}
}