bdk_tx 0.2.0

use alloc::sync::Arc;
use alloc::vec::Vec;
use core::fmt;

use bitcoin::{absolute, psbt, Amount, OutPoint, Sequence, Transaction, TxOut, Txid};
use miniscript::{bitcoin, plan::Plan};

use crate::{
    collections::{HashMap, HashSet},
    input::CoinbaseMismatch,
    ConfirmationStatus, FromPsbtInputError, Input, RbfSet,
};

/// Tx with confirmation status.
pub type TxWithStatus<T> = (T, Option<ConfirmationStatus>);

/// Our canonical view of unspent outputs.
#[derive(Debug, Clone)]
pub struct CanonicalUnspents {
    txs: HashMap<Txid, Arc<Transaction>>,
    statuses: HashMap<Txid, ConfirmationStatus>,
    spends: HashMap<OutPoint, Txid>,
}

impl CanonicalUnspents {
    /// Construct [`CanonicalUnspents`] from an iterator of txs with confirmation status.
    pub fn new<T>(canonical_txs: impl IntoIterator<Item = TxWithStatus<T>>) -> Self
    where
        T: Into<Arc<Transaction>>,
    {
        let mut txs = HashMap::new();
        let mut statuses = HashMap::new();
        let mut spends = HashMap::new();
        for (tx, status) in canonical_txs {
            let tx: Arc<Transaction> = tx.into();
            let txid = tx.compute_txid();
            spends.extend(tx.input.iter().map(|txin| (txin.previous_output, txid)));
            txs.insert(txid, tx);
            if let Some(status) = status {
                statuses.insert(txid, status);
            }
        }
        Self {
            txs,
            statuses,
            spends,
        }
    }

    /// Extract txs in the set of `replace` from the canonical view of unspents.
    ///
    /// Returns the [`RbfSet`] if the replacements are valid and succesfully extracted.
    /// Errors if the replacements cannot be extracted (e.g. due to missing data).
    pub fn extract_replacements(
        &mut self,
        replace: impl IntoIterator<Item = Txid>,
    ) -> Result<RbfSet, ExtractReplacementsError> {
        let mut rbf_txs = replace
            .into_iter()
            .map(|txid| -> Result<(Txid, Arc<Transaction>), _> {
                let tx = self
                    .txs
                    .get(&txid)
                    .cloned()
                    .ok_or(ExtractReplacementsError::TransactionNotFound(txid))?;
                if tx.is_coinbase() {
                    return Err(ExtractReplacementsError::CannotReplaceCoinbase);
                }
                Ok((tx.compute_txid(), tx))
            })
            .collect::<Result<HashMap<_, _>, _>>()?;

        // Find descendants of the original txs. Descendants are evicted by RBF and count toward
        // the required replacement fee.
        let mut descendants = HashMap::<Txid, Arc<Transaction>>::new();
        let mut visited = HashSet::<Txid>::new();
        let mut to_visit = rbf_txs
            .iter()
            .map(|(txid, tx)| (*txid, tx.clone()))
            .collect::<Vec<_>>();
        while let Some((txid, tx)) = to_visit.pop() {
            if !visited.insert(txid) {
                continue;
            }

            for vout in 0..tx.output.len() as u32 {
                let spent_outpoint = OutPoint::new(txid, vout);
                let Some(child_txid) = self.spends.get(&spent_outpoint).copied() else {
                    continue;
                };
                let Some(child_tx) = self.txs.get(&child_txid).cloned() else {
                    continue;
                };

                descendants
                    .entry(child_txid)
                    .or_insert_with(|| child_tx.clone());
                to_visit.push((child_txid, child_tx));
            }
        }

        for txid in descendants.keys() {
            rbf_txs.remove(txid);
        }

        // Find prev outputs of all txs in the set.
        // Fail when a prev output is not found. We need to use the prevouts to determine fee for RBF!
        let prev_txouts = rbf_txs
            .values()
            .chain(descendants.values())
            .flat_map(|tx| &tx.input)
            .map(|txin| txin.previous_output)
            .map(|op| -> Result<(OutPoint, TxOut), _> {
                let txout = self
                    .txs
                    .get(&op.txid)
                    .and_then(|tx| tx.output.get(op.vout as usize))
                    .cloned()
                    .ok_or(ExtractReplacementsError::PreviousOutputNotFound(op))?;
                Ok((op, txout))
            })
            .collect::<Result<HashMap<_, _>, _>>()?;

        let descendant_fee = descendants
            .values()
            .map(|tx| {
                let input_sum: Amount = tx
                    .input
                    .iter()
                    .map(|txin| prev_txouts[&txin.previous_output].value)
                    .sum();
                let output_sum: Amount = tx.output.iter().map(|txout| txout.value).sum();
                input_sum - output_sum
            })
            .sum::<Amount>();

        // Remove rbf txs (and their descendants) from canonical unspents.
        let to_remove_from_canonical_unspents = rbf_txs.keys().chain(descendants.keys());
        for txid in to_remove_from_canonical_unspents {
            if let Some(tx) = self.txs.remove(txid) {
                self.statuses.remove(txid);
                for txin in &tx.input {
                    self.spends.remove(&txin.previous_output);
                }
            }
        }

        let prev_txouts: HashMap<_, _> = rbf_txs
            .values()
            .flat_map(|tx| &tx.input)
            .map(|txin| txin.previous_output)
            .filter_map(|op| prev_txouts.get(&op).map(|txout| (op, txout.clone())))
            .collect();

        Ok(
            RbfSet::new(rbf_txs.into_values(), descendant_fee, prev_txouts)
                .expect("must not have missing prevouts"),
        )
    }

    /// Whether outpoint is a leaf (unspent).
    pub fn is_unspent(&self, outpoint: OutPoint) -> bool {
        if self.spends.contains_key(&outpoint) {
            return false;
        }
        match self.txs.get(&outpoint.txid) {
            Some(tx) => {
                let vout: usize = outpoint.vout.try_into().expect("vout must fit into usize");
                vout < tx.output.len()
            }
            None => false,
        }
    }

    /// Try get leaf (unspent) of given `outpoint`.
    pub fn try_get_unspent(&self, outpoint: OutPoint, plan: Plan) -> Option<Input> {
        if self.spends.contains_key(&outpoint) {
            return None;
        }
        let prev_tx = Arc::clone(self.txs.get(&outpoint.txid)?);
        Input::from_prev_tx(
            plan,
            prev_tx,
            outpoint.vout.try_into().expect("vout must fit into usize"),
            self.statuses.get(&outpoint.txid).cloned(),
        )
        .ok()
    }

    /// Try get leaves of given `outpoints`.
    pub fn try_get_unspents<'a, O>(&'a self, outpoints: O) -> impl Iterator<Item = Input> + 'a
    where
        O: IntoIterator<Item = (OutPoint, Plan)>,
        O::IntoIter: 'a,
    {
        outpoints
            .into_iter()
            .filter_map(|(op, plan)| self.try_get_unspent(op, plan))
    }

    /// Try get foreign leaf (unspent).
    pub fn try_get_foreign_unspent(
        &self,
        outpoint: OutPoint,
        sequence: Sequence,
        psbt_input: psbt::Input,
        satisfaction_weight: usize,
        is_coinbase: bool,
        absolute_timelock: Option<absolute::LockTime>,
    ) -> Result<Input, GetForeignUnspentError> {
        if !self.is_unspent(outpoint) {
            return Err(GetForeignUnspentError::OutputIsAlreadySpent(outpoint));
        }
        if let Some(prev_tx) = self.txs.get(&outpoint.txid) {
            let non_witness_utxo = psbt_input.non_witness_utxo.as_ref();
            if non_witness_utxo.is_some() && non_witness_utxo != Some(prev_tx) {
                return Err(GetForeignUnspentError::UtxoMismatch(outpoint));
            }
            let witness_utxo = psbt_input.witness_utxo.as_ref();
            if witness_utxo.is_some()
                && psbt_input.witness_utxo.as_ref() != prev_tx.output.get(outpoint.vout as usize)
            {
                return Err(GetForeignUnspentError::UtxoMismatch(outpoint));
            }
            if is_coinbase != prev_tx.is_coinbase() {
                return Err(GetForeignUnspentError::Coinbase(CoinbaseMismatch {
                    txid: outpoint.txid,
                    expected: is_coinbase,
                    got: prev_tx.is_coinbase(),
                }));
            }
        }
        let status = self.statuses.get(&outpoint.txid).cloned();
        Input::from_psbt_input(
            outpoint,
            sequence,
            psbt_input,
            satisfaction_weight,
            status,
            is_coinbase,
            absolute_timelock,
        )
        .map_err(GetForeignUnspentError::FromPsbtInput)
    }

    /// Try get foreign leaves (unspent).
    pub fn try_get_foreign_unspents<'a, O>(
        &'a self,
        outpoints: O,
    ) -> impl Iterator<Item = Result<Input, GetForeignUnspentError>> + 'a
    where
        O: IntoIterator<
            Item = (
                OutPoint,
                Sequence,
                psbt::Input,
                usize,
                bool,
                Option<absolute::LockTime>,
            ),
        >,
        O::IntoIter: 'a,
    {
        outpoints
            .into_iter()
            .map(|(op, seq, input, sat_wu, is_coinbase, absolute_timelock)| {
                self.try_get_foreign_unspent(op, seq, input, sat_wu, is_coinbase, absolute_timelock)
            })
    }
}

/// Canonical unspents error
#[derive(Debug)]
pub enum GetForeignUnspentError {
    /// Invalid parameter for `is_coinbase`
    Coinbase(CoinbaseMismatch),
    /// Error creating an input from a PSBT input
    FromPsbtInput(FromPsbtInputError),
    /// Cannot get unspent input from output that is already spent
    OutputIsAlreadySpent(OutPoint),
    /// The witness or non-witness UTXO in the PSBT input does not match the expected outpoint
    UtxoMismatch(OutPoint),
}

impl fmt::Display for GetForeignUnspentError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Coinbase(err) => write!(f, "{err}"),
            Self::FromPsbtInput(err) => write!(f, "{err}"),
            Self::OutputIsAlreadySpent(op) => {
                write!(f, "outpoint is already spent: {op}")
            }
            Self::UtxoMismatch(op) => write!(f, "UTXO mismatch: {op}"),
        }
    }
}

#[cfg(feature = "std")]
impl std::error::Error for GetForeignUnspentError {}

/// Error when attempting to do [`extract_replacements`](CanonicalUnspents::extract_replacements).
#[derive(Debug)]
pub enum ExtractReplacementsError {
    /// Transaction not found in canonical unspents
    TransactionNotFound(Txid),
    /// Cannot replace a coinbase transaction
    CannotReplaceCoinbase,
    /// Previous output not found for input
    PreviousOutputNotFound(OutPoint),
}

impl fmt::Display for ExtractReplacementsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::TransactionNotFound(txid) => write!(f, "transaction not found: {txid}"),
            Self::CannotReplaceCoinbase => write!(f, "cannot replace a coinbase transaction"),
            Self::PreviousOutputNotFound(op) => write!(f, "previous output not found: {op}"),
        }
    }
}

#[cfg(feature = "std")]
impl std::error::Error for ExtractReplacementsError {}

#[cfg_attr(coverage_nightly, coverage(off))]
#[cfg(test)]
mod tests {
    use super::*;
    use bitcoin::{absolute, transaction, Amount, ScriptBuf, TxIn};

    fn funding_tx(output_values: &[u64]) -> Transaction {
        Transaction {
            version: transaction::Version::TWO,
            lock_time: absolute::LockTime::ZERO,
            input: Vec::new(),
            output: output_values
                .iter()
                .map(|value| TxOut {
                    value: Amount::from_sat(*value),
                    script_pubkey: ScriptBuf::new(),
                })
                .collect(),
        }
    }

    fn tx_spending(inputs: &[OutPoint], output_values: &[u64]) -> Transaction {
        Transaction {
            version: transaction::Version::TWO,
            lock_time: absolute::LockTime::ZERO,
            input: inputs
                .iter()
                .map(|previous_output| TxIn {
                    previous_output: *previous_output,
                    ..Default::default()
                })
                .collect(),
            output: output_values
                .iter()
                .map(|value| TxOut {
                    value: Amount::from_sat(*value),
                    script_pubkey: ScriptBuf::new(),
                })
                .collect(),
        }
    }

    fn prevout(tx: &Transaction, vout: u32) -> OutPoint {
        OutPoint::new(tx.compute_txid(), vout)
    }

    fn unconfirmed_txs(txs: Vec<Transaction>) -> Vec<TxWithStatus<Transaction>> {
        txs.into_iter().map(|tx| (tx, None)).collect()
    }

    /// Counts fees from the original tx and its full descendant chain.
    ///
    /// Fee floor: parent 1_000 + child 2_000 + grandchild 3_000 = 6_000 sats.
    #[test]
    fn test_extract_replacements_counts_transitive_descendant_fees() {
        let funding = funding_tx(&[50_000]);
        let parent = tx_spending(&[prevout(&funding, 0)], &[49_000]);
        let child = tx_spending(&[prevout(&parent, 0)], &[47_000]);
        let grandchild = tx_spending(&[prevout(&child, 0)], &[44_000]);
        let parent_txid = parent.compute_txid();
        let parent_feerate = Amount::from_sat(1_000) / parent.weight();
        let mut canonical_unspents =
            CanonicalUnspents::new(unconfirmed_txs(vec![funding, parent, child, grandchild]));

        let rbf_set = canonical_unspents
            .extract_replacements([parent_txid])
            .expect("replacement set should extract");

        let txids = rbf_set.txids().collect::<Vec<_>>();
        assert_eq!(txids, vec![parent_txid]);
        let rbf_params = rbf_set.selector_rbf_params();
        assert_eq!(rbf_params.to_cs_replace().fee, 6_000);
        assert_eq!(rbf_params.max_feerate(), parent_feerate);
    }

    /// Keeps requested descendants out of original txids while still fee-counting them.
    #[test]
    fn test_extract_replacements_moves_requested_descendant_out_of_original_txids() {
        let funding = funding_tx(&[50_000]);
        let parent = tx_spending(&[prevout(&funding, 0)], &[49_000]);
        let child = tx_spending(&[prevout(&parent, 0)], &[47_000]);
        let parent_txid = parent.compute_txid();
        let child_txid = child.compute_txid();
        let mut canonical_unspents =
            CanonicalUnspents::new(unconfirmed_txs(vec![funding, parent, child]));

        let rbf_set = canonical_unspents
            .extract_replacements([parent_txid, child_txid])
            .expect("replacement set should extract");

        let txids = rbf_set.txids().collect::<Vec<_>>();
        assert_eq!(txids.len(), 1);
        assert!(txids.contains(&parent_txid));
        assert!(!txids.contains(&child_txid));
        assert_eq!(rbf_set.selector_rbf_params().to_cs_replace().fee, 3_000);
    }
}