bark-wallet 0.1.1

use ark::vtxo::policy::signing::VtxoSigner;
use log::{debug, error, info, trace, warn};

use bitcoin_ext::{BlockDelta, P2TR_DUST, TxStatus};
use crate::exit::models::{
	ExitError, ExitAwaitingDeltaState, ExitProcessingState, ExitClaimInProgressState, ExitClaimableState,
	ExitClaimedState, ExitState, ExitStartState, ExitTx, ExitTxOrigin, ExitTxStatus,
};
use crate::exit::progress::{ExitProgressError, ExitStateProgress, ProgressContext};
use crate::exit::progress::util::{count_broadcast, count_confirmed, estimate_exit_cost};
use crate::onchain::ExitUnilaterally;

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl ExitStateProgress for ExitState {
	async fn progress(
		self,
		ctx: &mut ProgressContext<'_>,
		onchain: &mut dyn ExitUnilaterally,
	) -> anyhow::Result<ExitState, ExitProgressError> {
		match self {
			ExitState::Start(s) => s.progress(ctx, onchain).await,
			ExitState::Processing(s) => s.progress(ctx, onchain).await,
			ExitState::AwaitingDelta(s) => s.progress(ctx, onchain).await,
			ExitState::Claimable(s) => s.progress(ctx, onchain).await,
			ExitState::ClaimInProgress(s) => s.progress(ctx, onchain).await,
			ExitState::Claimed(s) => s.progress(ctx, onchain).await,
		}
	}
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl ExitStateProgress for ExitStartState {
	async fn progress(
		self,
		ctx: &mut ProgressContext<'_>,
		onchain: &mut dyn ExitUnilaterally,
	) -> anyhow::Result<ExitState, ExitProgressError> {
		let id = ctx.vtxo.id();
		info!("Checking if VTXO can be exited: {}", id);

		// Ensure the VTXO has a valid amount
		if ctx.vtxo.amount() < P2TR_DUST {
			return Err(ExitError::DustLimit { vtxo: ctx.vtxo.amount(), dust: P2TR_DUST }.into());
		}

		// Ensure we can afford to exit this VTXO
		let total_fee = estimate_exit_cost([ctx.vtxo], ctx.fee_rate);
		let balance = onchain.get_balance();
		if balance < total_fee {
			return Err(ExitError::InsufficientFeeToStart {
				balance,
				total_fee,
				fee_rate: ctx.fee_rate
			}.into());
		}
		info!("Validated VTXO {}, exit process can now begin", id);

		Ok(ExitState::new_processing(
			ctx.wallet.chain.tip().await.unwrap_or(self.tip_height),
			ctx.exit_txids.iter().cloned(),
		))
	}
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl ExitStateProgress for ExitProcessingState {
	async fn progress(
		self,
		ctx: &mut ProgressContext<'_>,
		onchain: &mut dyn ExitUnilaterally,
	) -> anyhow::Result<ExitState, ExitProgressError> {
		assert_eq!(self.transactions.len(), ctx.exit_txids.len());

		let tip = ctx.tip_height().await?;
		let mut transactions = self.transactions.clone();

		for i in 0..transactions.len() {
			match progress_exit_tx(&transactions[i], ctx, onchain).await {
				Ok(status) => transactions[i].status = status,
				Err(e) => {
					// We may need to commit any changes we have
					if self.transactions != transactions {
						let state = ExitState::new_processing_from_transactions(tip, transactions);
						return Err(ExitProgressError {
							state: Some(state),
							error: e,
						});
					}
					return Err(e.into());
				},
			}
		}

		// Report the current status to the user
		let prev_confirmed = count_confirmed(&self.transactions);
		let now_confirmed = count_confirmed(&transactions);
		if now_confirmed == transactions.len() {
			info!("Exit for VTXO ({}) has been fully confirmed, waiting for funds to become \
				spendable...", ctx.vtxo.id(),
			);
			let conf_block = transactions
				.iter()
				.filter_map(|exit| exit.status.confirmed_in())
				.max_by(|a, b| a.height.cmp(&b.height))
				.unwrap();

			let clause = ctx.wallet.find_signable_clause(ctx.vtxo).await
				.ok_or_else(|| ExitError::ClaimMissingSignableClause { vtxo: ctx.vtxo.id() })?;

			let wait_delta = clause.sequence().map_or(0, |csv| csv.0) as BlockDelta;
			return Ok(ExitState::new_awaiting_delta(tip, *conf_block, wait_delta));
		}
		if now_confirmed != prev_confirmed {
			info!("Exit for VTXO ({}) now has {} confirmed transactions with {} more required.",
				ctx.vtxo.id(), now_confirmed, transactions.len() - now_confirmed,
			);
		} else {
			let prev_broadcast = count_broadcast(&self.transactions);
			let now_broadcast = count_broadcast(&transactions);
			if now_broadcast == transactions.len() {
				info!("Exit for VTXO ({}) has been fully broadcast, waiting for {} transactions \
					to confirm...", ctx.vtxo.id(), now_confirmed,
				);
			} else if prev_broadcast != now_broadcast {
				let remaining = transactions.len() - now_broadcast;
				if prev_broadcast > now_broadcast {
					warn!("An exit transaction for VTXO ({}) appears to have fallen out of the \
						mempool", ctx.vtxo.id(),
					);
				}
				info!("Exit for VTXO ({}) now has {} broadcast transactions with {} more required.",
					ctx.vtxo.id(), now_broadcast, remaining,
				);
			}
		}

		if self.transactions != transactions {
			debug!("VTXO exit transactions updated: {:?}", transactions);
			Ok(ExitState::new_processing_from_transactions(tip, transactions))
		} else {
			Ok(self.into())
		}
	}
}

async fn progress_exit_tx(
	exit: &ExitTx,
	ctx: &mut ProgressContext<'_>,
	onchain: &mut dyn ExitUnilaterally,
) -> anyhow::Result<ExitTxStatus, ExitError> {
	match &exit.status {
		ExitTxStatus::VerifyInputs => {
			debug!("Verifying inputs for exit tx {}", exit.txid);
			let inputs = ctx.get_unique_inputs(exit.txid).await?;
			ctx.check_status_from_inputs(exit, &inputs).await
		},
		ExitTxStatus::AwaitingInputConfirmation { txids } => {
			debug!("Checking if the {} remaining inputs for exit tx {} have confirmed",
				txids.len(), exit.txid,
			);
			ctx.check_status_from_inputs(exit, &txids).await
		}
		ExitTxStatus::NeedsSignedPackage => {
			// Before attempting to create a package, we should verify another party hasn't
			// already broadcast this transaction
			let new_status = ctx.get_exit_tx_status(exit).await?;
			if matches!(new_status, ExitTxStatus::NeedsSignedPackage) {
				debug!("Creating exit package for exit tx {}", exit.txid);
				let child_tx = {
					let package = ctx.tx_manager.get_package(exit.txid)?;
					let guard = package.read().await;
					assert_eq!(guard.child, None);

					ctx.create_exit_cpfp_tx(&guard.exit.tx, onchain, None)?
				};

				// Update the transaction manager so our package can be broadcast later
				let origin = ExitTxOrigin::Wallet { confirmed_in: None };
				let child_txid = ctx.tx_manager.set_wallet_child_tx(
					exit.txid, child_tx, origin,
				).await?;

				debug!("CPFP created with txid {} for exit tx {}", child_txid, exit.txid);
				Ok(ExitTxStatus::NeedsBroadcasting { child_txid, origin })
			} else {
				debug!("Exit tx {} has likely been broadcast by another party", exit.txid);
				Ok(new_status)
			}
		}
		ExitTxStatus::NeedsReplacementPackage { .. } => {
			// Ensure we still need to replace the package
			match ctx.get_exit_tx_status(exit).await? {
				ExitTxStatus::NeedsReplacementPackage { min_fee_rate, min_fee } => {
					debug!("Creating replacement exit package with a fee rate of at least \
						{}sats/kWu and a minimum fee of {} for exit tx {}",
						min_fee_rate, min_fee, exit.txid,
					);
					let child_tx = ctx.create_exit_cpfp_tx(
						&ctx.tx_manager.get_package(exit.txid)?.read().await.exit.tx,
						onchain,
						Some((min_fee_rate, min_fee)),
					)?;

					// Update the transaction manager so our package can be broadcast later
					let origin = ExitTxOrigin::Wallet { confirmed_in: None };
					let child_txid = ctx.tx_manager.set_wallet_child_tx(
						exit.txid, child_tx, origin,
					).await?;

					debug!("RBF CPFP created with txid {} for exit tx {}", child_txid, exit.txid);
					Ok(ExitTxStatus::NeedsBroadcasting { child_txid, origin })
				},
				s => {
					debug!("Status has changed for exit tx {}, no longer creating a replacement \
						package", exit.txid,
					);
					Ok(s)
				},
			}
		},
		ExitTxStatus::NeedsBroadcasting { child_txid, .. } => {
			debug!("Checking if exit tx {} has been broadcast with CPFP tx {}",
				exit.txid, child_txid,
			);
			let status = ctx.get_exit_child_status(&exit, *child_txid).await?;
			match status {
				ExitTxStatus::NeedsBroadcasting { child_txid: new_child_txid, .. } => {
					if new_child_txid != *child_txid {
						warn!("Exit tx {} has a different child txid. Expected: {} Found: {}",
							exit.txid, child_txid, new_child_txid,
						);
					}
					debug!("Attempting to broadcast exit tx {} with child tx {}",
						exit.txid, child_txid,
					);
					let package = ctx.tx_manager.get_package(exit.txid)?;
					let guard = package.read().await;
					let status = ctx.tx_manager.broadcast_package(&*guard).await?;
					if matches!(status, TxStatus::Mempool) {
						debug!("Commiting exit CPFP {} to database", new_child_txid);
						let tx = &guard.child.as_ref().expect("child can't be missing").info.tx;
						onchain.store_signed_p2a_cpfp(tx).await
							.map_err(|e| ExitError::ExitPackageStoreFailure {
								txid: exit.txid,
								error: e.to_string(),
							})?;
					}

					// Finally, we can go to the next state
					ctx.get_exit_child_status(&exit, new_child_txid).await
				},
				_ => {
					debug!("Exit tx {} needed broadcasting but has changed status to: {}",
						exit.txid, status,
					);
					Ok(status)
				},
			}
		},
		ExitTxStatus::BroadcastWithCpfp { child_txid, .. } => {
			let new_status = ctx.get_exit_child_status(exit, *child_txid).await?;
			match new_status {
				ExitTxStatus::Confirmed { block, .. } => {
					debug!("Exit tx {} confirmed at height {}", exit.txid, block.height);
				}
				_ => {}
			}
			Ok(new_status)
		},
		ExitTxStatus::Confirmed { child_txid, block, .. } => {
			// Handle cases where we might get a block-reorg so our transaction may unconfirm
			let new_status = ctx.get_exit_child_status(exit, *child_txid).await?;
			match &new_status {
				ExitTxStatus::Confirmed { child_txid: new_txid, block: new_block, .. } => {
					if new_block != block || new_txid != child_txid {
						warn!("Exit transaction {} was confirmed with block {} but it has been \
							replaced by {} in block {}",
							exit.txid, block.hash, new_txid, new_block.hash
						);
					}
				},
				_ => {
					warn!("Exit transaction {} was confirmed at height {} but it's now unconfirmed",
						exit.txid, block.height
					);
				},
			}
			Ok(new_status)
		}
	}
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl ExitStateProgress for ExitAwaitingDeltaState {
	async fn progress(
		self,
		ctx: &mut ProgressContext<'_>,
		_onchain: &mut dyn ExitUnilaterally,
	) -> anyhow::Result<ExitState, ExitProgressError> {
		let tip = ctx.tip_height().await?;

		// Ensure the exit transaction hasn't disappeared from the mempool due to a reorg
		if !ctx.check_confirmed(ctx.vtxo.point().txid).await {
			error!("Exit for VTXO ({}) is no longer confirmed, verifying all transactions...",
				ctx.vtxo.id(),
			);
			return Ok(ExitState::new_processing(
				tip, ctx.exit_txids.iter().cloned(),
			));
		}

		// Inform the user of any progress
		if tip >= self.claimable_height {
			info!("Exit for VTXO ({}) is spendable!", ctx.vtxo.id());
			let spendable_block = ctx.get_block_ref(self.claimable_height).await?;
			Ok(ExitState::new_claimable(tip, spendable_block, None))
		} else {
			info!("Waiting for {} more confirmations until exit for VTXO ({}) is spendable...",
				self.claimable_height - tip, ctx.vtxo.id(),
			);
			Ok(self.into())
		}
	}
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl ExitStateProgress for ExitClaimableState {
	async fn progress(
		self,
		ctx: &mut ProgressContext<'_>,
		_onchain: &mut dyn ExitUnilaterally,
	) -> anyhow::Result<ExitState, ExitProgressError> {
		let tip = ctx.tip_height().await?;

		// We should verify the current block hasn't been reorganized.
		let spendable_block = ctx.get_block_ref(self.claimable_since.height).await?;
		if spendable_block.hash != self.claimable_since.hash {
			return Ok(ExitState::new_claimable(tip, spendable_block, None));
		}

		// We can avoid scanning the whole chain provided there hasn't been a re-org
		let scan_height = if let Some(block) = &self.last_scanned_block {
			// Double check we haven't had a re-org since the last scan
			if ctx.get_block_ref(block.height).await?.hash == block.hash {
				block.height
			} else {
				self.claimable_since.height
			}
		} else {
			self.claimable_since.height
		};

		// Check if the VTXO exit has been spent
		let point = ctx.vtxo.point();
		let result = ctx.wallet.chain
			.txs_spending_inputs(
				vec![point],
				scan_height,
			).await
			.map_err(|e| ExitError::TransactionRetrievalFailure {
				txid: ctx.vtxo.point().txid, error: e.to_string(),
			})?;

		if let Some((txid, status)) = result.get(&point) {
			match status {
				TxStatus::Confirmed(block) => {
					debug!("Tx {} has successfully claimed VTXO {}", txid, ctx.vtxo.id());
					Ok(ExitState::new_claimed(tip, txid.clone(), *block))
				},
				TxStatus::Mempool => {
					debug!("Tx {} is attempting to claim VTXO {}", txid, ctx.vtxo.id());
					Ok(ExitState::new_claim_in_progress(tip, self.claimable_since, txid.clone()))
				},
				TxStatus::NotFound => unreachable!(),
			}
		} else {
			// Make sure the wallet is aware of the exit
			debug!("VTXO is still spendable: {}", ctx.vtxo.id());
			let tip_block = Some(ctx.get_block_ref(tip).await?);
			Ok(ExitState::new_claimable(tip, self.claimable_since, tip_block))
		}
	}
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl ExitStateProgress for ExitClaimInProgressState {
	async fn progress(
		self,
		ctx: &mut ProgressContext<'_>,
		_onchain: &mut dyn ExitUnilaterally,
	) -> anyhow::Result<ExitState, ExitProgressError> {
		// Wait for confirmation of the spending transaction
		let tip = ctx.tip_height().await?;
		match ctx.tx_manager.tx_status(self.claim_txid).await? {
			TxStatus::Confirmed(block) => {
				debug!("Tx {} has successfully spent VTXO {}", self.claim_txid, ctx.vtxo.id());
				Ok(ExitState::new_claimed(tip, self.claim_txid, block))
			},
			TxStatus::Mempool => {
				trace!("Still waiting for TX {} to be confirmed", self.claim_txid);
				Ok(self.into())
			},
			TxStatus::NotFound => {
				warn!("TX {} has dropped from the mempool, VTXO {} is spendable again",
					self.claim_txid, ctx.vtxo.id(),
				);
				Ok(ExitState::new_claimable(tip, self.claimable_since, None))
			},
		}
	}
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl ExitStateProgress for ExitClaimedState {
	async fn progress(
		self,
		ctx: &mut ProgressContext<'_>,
		_onchain: &mut dyn ExitUnilaterally,
	) -> anyhow::Result<ExitState, ExitProgressError> {
		trace!("Exit for VTXO {} is spent!", ctx.vtxo.id());
		Ok(self.into())
	}
}