rostrum 14.0.1

use crate::chaindef::OutPoint;
use crate::chaindef::{OutPointHash, Transaction};
use crate::daemon::Daemon;
use crate::index::index_transaction;
use crate::metrics::{self, Metrics};
use crate::query::tx::TxQuery;
use crate::store::{DBContents, DBStore};
use crate::ttorsort::ttor_sorted_mempool_entries;
use crate::writebatch::{SpentUpdateBatch, WriteBatch};
use anyhow::{Context, Result};
use bitcoin_hashes::hex::ToHex;
use bitcoincash::hash_types::Txid;
use futures_util::{stream, StreamExt};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock};

const VSIZE_BIN_WIDTH: u32 = 100_000; // in vbytes

/// Fake height value used to signify that a transaction is in the memory pool.
pub const MEMPOOL_HEIGHT: u32 = 0x7FFF_FFFF;

pub enum ConfirmationState {
    Indeterminate,
    Confirmed,
    InMempool,
    UnconfirmedParent,
}

pub struct MempoolEntry {
    fee: u64,    // in satoshis
    size: usize, // in bytes
    fee_per_byte: f32,
    tx: Transaction,
}

impl MempoolEntry {
    pub(crate) fn new(fee: u64, tx: Transaction) -> MempoolEntry {
        let size = tx.size();
        MempoolEntry {
            fee,
            size,
            fee_per_byte: fee as f32 / size as f32,
            tx,
        }
    }

    pub fn fee(&self) -> u64 {
        self.fee
    }

    pub fn size(&self) -> usize {
        self.size
    }

    pub fn fee_per_byte(&self) -> f32 {
        self.fee_per_byte
    }

    pub fn tx(&self) -> &Transaction {
        &self.tx
    }
}

struct Stats {
    count: prometheus::IntGauge,
    update: prometheus::HistogramVec,
    vsize: prometheus::GaugeVec,
    max_fee_rate: Mutex<f32>,
}

impl Stats {
    fn start_timer(&self, step: &str) -> prometheus::HistogramTimer {
        self.update.with_label_values(&[step]).start_timer()
    }

    async fn update(&self, entries: &[&MempoolEntry]) {
        let mut bands: Vec<(f32, u32)> = vec![];
        let mut fee_rate = 1.0f32; // [sat/vbyte]
        let mut vsize = 0u32; // vsize of transactions paying <= fee_rate
        for e in entries {
            while fee_rate < e.fee_per_byte() {
                bands.push((fee_rate, vsize));
                fee_rate *= 2.0;
            }
            vsize += e.size() as u32;
        }
        let mut max_fee_rate = self.max_fee_rate.lock().await;
        loop {
            bands.push((fee_rate, vsize));
            if fee_rate < *max_fee_rate {
                fee_rate *= 2.0;
                continue;
            }
            *max_fee_rate = fee_rate;
            break;
        }
        drop(max_fee_rate);
        for (fee_rate, vsize) in bands {
            // labels should be ordered by fee_rate value
            let label = format!("≤{:10.0}", fee_rate);
            self.vsize
                .with_label_values(&[&label])
                .set(f64::from(vsize));
        }
    }
}

pub struct Tracker {
    items: RwLock<HashMap<Txid, MempoolEntry>>,
    index: Arc<DBStore>,
    histogram: Mutex<Vec<(f32, u32)>>,
    stats: Stats,
}

impl Tracker {
    pub fn new(mempool_path: &std::path::Path, metrics: &Metrics) -> Tracker {
        let store = {
            // start with an empty mempool
            DBStore::destroy(mempool_path);
            DBStore::open(DBContents::MempoolIndex, mempool_path, metrics, false)
                .expect("creating mempool db faled")
        };
        Tracker {
            items: RwLock::new(HashMap::new()),
            index: Arc::new(store),
            histogram: Mutex::new(vec![]),
            stats: Stats {
                count: metrics.gauge_int(prometheus::Opts::new(
                    "rostrum_mempool_count",
                    "# of mempool transactions",
                )),
                update: metrics.histogram_vec(
                    "rostrum_mempool_update",
                    "Time to update mempool (in seconds)",
                    &["step"],
                    metrics::default_duration_buckets(),
                ),
                vsize: metrics.gauge_float_vec(
                    prometheus::Opts::new(
                        "rostrum_mempool_vsize",
                        "Total vsize of transactions paying at most given fee rate",
                    ),
                    &["fee_rate"],
                ),
                max_fee_rate: Mutex::new(1.0),
            },
        }
    }

    pub async fn get_txn(&self, txid: &Txid) -> Option<Transaction> {
        self.items
            .read()
            .await
            .get(txid)
            .map(|stats| stats.tx.clone())
    }

    pub async fn get_all_txn(&self) -> Vec<Txid> {
        self.items.read().await.keys().copied().collect()
    }

    pub async fn get_txn_partially_matching_scriptsig(&self, scriptsig: &[u8]) -> Vec<Txid> {
        self.items
            .read()
            .await
            .iter()
            .filter_map(|(txid, entry)| {
                if entry.tx.input.iter().any(|i| {
                    i.script_sig
                        .to_bytes()
                        .windows(scriptsig.len())
                        .any(|window| window == scriptsig)
                }) {
                    Some(*txid)
                } else {
                    None
                }
            })
            .collect()
    }

    pub async fn get_txn_partially_matching_scriptpubkey(&self, scriptpubkey: &[u8]) -> Vec<Txid> {
        self.items
            .read()
            .await
            .iter()
            .filter_map(|(txid, entry)| {
                if entry.tx.output.iter().any(|o| {
                    o.script_pubkey
                        .to_bytes()
                        .windows(scriptpubkey.len())
                        .any(|window| window == scriptpubkey)
                }) {
                    Some(*txid)
                } else {
                    None
                }
            })
            .collect()
    }

    pub async fn has_txn(&self, txid: &Txid) -> bool {
        self.items.read().await.contains_key(txid)
    }

    pub async fn get_fee(&self, txid: &Txid) -> Option<u64> {
        self.items.read().await.get(txid).map(|stats| stats.fee())
    }

    pub async fn get_count(&self) -> usize {
        self.items.read().await.len()
    }

    #[cfg(bch)]
    pub async fn contains_funding_tx(&self, prevout: &OutPoint) -> bool {
        self.items.read().await.contains_key(&prevout.txid)
    }
    #[cfg(nexa)]
    /// Transaction funding the given utxo (represented as OutPoint)
    pub async fn contains_funding_tx(&self, prevout: &OutPoint) -> bool {
        use bitcoin_hashes::Hash;

        use crate::encode::outpoint_hash;
        use crate::indexes::outputindex::OutputIndexRow;
        use crate::indexes::DBRow;

        let key = outpoint_hash(prevout);
        self.index
            .exists_32bit_key(OutputIndexRow::CF, key.into_inner())
            .await
    }
    /// Returns vector of (fee_rate, vsize) pairs, where fee_{n-1} > fee_n and vsize_n is the
    /// total virtual size of mempool transactions with fee in the bin [fee_{n-1}, fee_n].
    /// Note: fee_{-1} is implied to be infinite.
    pub async fn fee_histogram(&self) -> Vec<(f32, u32)> {
        self.histogram.lock().await.clone()
    }

    pub fn index(&self) -> &Arc<DBStore> {
        &self.index
    }

    async fn find_mempool_difference_daemon(
        &self,
        daemon: &Daemon,
    ) -> Result<(HashSet<Txid>, HashSet<Txid>)> {
        let daemon_txids = daemon
            .getmempooltxids()
            .await
            .context("failed to update mempool from daemon")?;

        let our: HashSet<Txid> = self.items.read().await.keys().cloned().collect();

        let new_txs = daemon_txids.difference(&our).cloned().collect();
        let removed_txs = our.difference(&daemon_txids).cloned().collect();
        Ok((new_txs, removed_txs))
    }

    /**
     * Called when we have a transaction and can assume we have its parents in mempool already.
     */
    pub async fn update_from_tx(&self, tx: Transaction, txquery: &TxQuery) -> Result<()> {
        let dummy_cache: HashMap<OutPointHash, u64> = HashMap::default();

        let txid = tx.txid();

        if self.items.read().await.contains_key(&txid) {
            return Ok(());
        }

        let entry = txquery.mempool_entry_from(self, tx, &dummy_cache).await?;
        self.add(&txid, entry, txquery).await;

        Ok(())
    }

    /**
     * Helper function to add transactions to mempool in TTOR-sorted order.
     * Collects entries, sorts them, and adds them.
     */
    async fn add_transactions_ttor_sorted(
        &self,
        new_txs: HashSet<Txid>,
        txquery: &TxQuery,
        error_prefix: &str,
    ) -> HashSet<Txid> {
        if new_txs.is_empty() {
            return HashSet::new();
        }

        let input_fee_cache: Arc<HashMap<OutPointHash, u64>> =
            Arc::new(txquery.create_input_cache(&new_txs).await);

        // Collect all mempool entries first
        let entries: Vec<MempoolEntry> = stream::iter(new_txs)
            .filter_map(|txid| {
                let cache = input_fee_cache.clone();

                async move {
                    match txquery.get_mempool_entry(self, &txid, &cache).await {
                        Ok(entry) => {
                            assert_eq!(entry.tx.txid(), txid);
                            Some(entry)
                        }
                        Err(err) => {
                            debug!("failed to get {} {}: {}", error_prefix, txid, err);
                            None
                        }
                    }
                }
            })
            .collect()
            .await;

        // Sort entries by TTOR before adding to mempool
        let sorted_entries = ttor_sorted_mempool_entries(entries);

        // Add entries in TTOR-sorted order
        let mut added: HashSet<Txid> = HashSet::new();
        for entry in sorted_entries {
            let txid = entry.tx.txid();
            self.add(&txid, entry, txquery).await;
            added.insert(txid);
        }

        added
    }

    /**
     * Add any new transaction batch provided via p2p network to our mempool.
     */
    pub async fn update_from_txids(
        &self,
        tx_invs: HashSet<Txid>,
        txquery: &TxQuery,
    ) -> HashSet<Txid> {
        // Filtering instead of calling `difference`, as tx_invs is likely to be small.

        let new_txs: HashSet<_> = stream::iter(tx_invs)
            .filter(|tx| {
                let tx = *tx;
                async move { !self.items.read().await.contains_key(&tx) }
            })
            .collect()
            .await;

        let added = self
            .add_transactions_ttor_sorted(new_txs, txquery, "p2p mempool entry")
            .await;

        self.update_fee_histogram().await;
        self.stats.count.set(self.items.read().await.len() as i64);
        added
    }

    /**
     * Find full difference in mempool, the daemon is the master, so any transactions
     * they don't have, but we have, we'll remove.
     */
    pub async fn update_from_daemon(
        &self,
        daemon: &Daemon,
        txquery: &TxQuery,
    ) -> Result<HashSet<Txid>> {
        // set of transactions where a change has occurred (either new or removed)

        let timer = self.stats.start_timer("fetch");
        let (new_txs, removed_txs) = self.find_mempool_difference_daemon(daemon).await?;
        timer.observe_duration();

        if new_txs.is_empty() && removed_txs.is_empty() {
            return Ok(HashSet::new());
        }

        let timer = self.stats.start_timer("add");

        let mut changed_txs = self
            .add_transactions_ttor_sorted(new_txs, txquery, "mempool entry")
            .await;

        timer.observe_duration();

        let timer = self.stats.start_timer("remove");

        // Collect transactions for removed_txs and sort by reverse TTOR (children before parents)
        let mut removed_transactions: Vec<Transaction> = Vec::new();
        for txid in &removed_txs {
            if let Some(tx) = self.get_txn(txid).await {
                removed_transactions.push(tx);
            }
        }

        // Sort by TTOR, then reverse to get children before parents
        let sorted_removed = ttor_sorted_mempool_entries(
            removed_transactions
                .into_iter()
                .map(|tx| MempoolEntry::new(0, tx)) // Fee doesn't matter for removal
                .collect(),
        );
        let reverse_sorted: Vec<_> = sorted_removed.into_iter().rev().collect();

        // Remove in reverse TTOR order (children before parents)
        for entry in reverse_sorted {
            let txid = entry.tx.txid();
            self.remove(&txid, txquery).await;
            changed_txs.insert(txid);
        }
        timer.observe_duration();

        let timer = self.stats.start_timer("fees");
        self.update_fee_histogram().await;
        timer.observe_duration();

        self.stats.count.set(self.items.read().await.len() as i64);
        Ok(changed_txs)
    }

    pub async fn tx_confirmation_state(
        &self,
        txid: &Txid,
        height: Option<u32>,
    ) -> ConfirmationState {
        if let Some(height) = height {
            if height != MEMPOOL_HEIGHT {
                return ConfirmationState::Confirmed;
            }
        }

        // Check if any of our inputs are unconfirmed
        if let Some(tx) = self.get_txn(txid).await {
            for input in tx.input.iter() {
                if self.contains_funding_tx(&input.previous_output).await {
                    return ConfirmationState::UnconfirmedParent;
                }
            }
            ConfirmationState::InMempool
        } else {
            // We don't see it in our mempool.
            ConfirmationState::Indeterminate
        }
    }

    async fn add(&self, txid: &Txid, entry: MempoolEntry, txquery: &TxQuery) {
        let writebatch = WriteBatch::new();

        // Transactions need to be added in ttor manner (no orphans); we require this for consistency.
        // We relay on create_prefilled_spent_data to fail if the transaction is an orphan.
        let prefilled_spent_data = match txquery.create_prefilled_spent_data(&entry.tx).await {
            Ok(data) => data,
            Err(e) => {
                debug!(
                    "mempool add: failed to create prefilled spent data for tx (orphan?) {}: {}",
                    txid.to_hex(),
                    e
                );
                return;
            }
        };
        let spentupdatebatch =
            SpentUpdateBatch::new().add_prefilled_spent_data(prefilled_spent_data.clone());

        index_transaction(
            &writebatch,
            &spentupdatebatch,
            &entry.tx,
            MEMPOOL_HEIGHT as usize,
        );

        self.index.write_batch(&writebatch);
        self.index.update_spends(&spentupdatebatch);

        // update our mempool after index is up-to-date
        {
            let mut items = self.items.write().await;
            items.insert(*txid, entry);
        }
    }

    async fn remove(&self, txid: &Txid, txquery: &TxQuery) {
        // Get transaction data first (with read lock, minimal time)
        let tx = {
            let items = self.items.read().await;
            match items.get(txid) {
                Some(entry) => entry.tx.clone(),
                None => {
                    // Already removed
                    return;
                }
            }
        };

        // Pre-fetch all transactions we need BEFORE acquiring write lock
        let writebatch = WriteBatch::new();
        let spentupdatebatch = SpentUpdateBatch::new().set_erase_only();

        let spentupdatebatch = match txquery.create_prefilled_spent_data(&tx).await {
            Ok(data) => spentupdatebatch.add_prefilled_spent_data(data),
            Err(e) => {
                debug!(
                    "mempool remove: failed to create prefilled spent data for tx (orphan?) {}: {}",
                    txid.to_hex(),
                    e
                );
                // erase what we can from mempool (maybe clear all mempool?)
                spentupdatebatch
            }
        };

        index_transaction(&writebatch, &spentupdatebatch, &tx, MEMPOOL_HEIGHT as usize);

        // Update database indexes BEFORE removing from mempool items
        self.index.erase_batch(&writebatch);
        self.index.update_spends(&spentupdatebatch);

        // Now remove from mempool items (index is already up-to-date)
        {
            let mut items = self.items.write().await;
            items.remove(txid);
        }
    }

    async fn update_fee_histogram(&self) {
        let items = self.items.read().await;
        let mut histogram = self.histogram.lock().await;

        let mut entries: Vec<&MempoolEntry> = items.values().collect();
        entries
            .sort_unstable_by(|e1, e2| e1.fee_per_byte().partial_cmp(&e2.fee_per_byte()).unwrap());
        *histogram = electrum_fees(&entries);
        self.stats.update(&entries).await;
    }
}

fn electrum_fees(entries: &[&MempoolEntry]) -> Vec<(f32, u32)> {
    let mut histogram = vec![];
    let mut bin_size = 0;
    let mut last_fee_rate = 0.0;
    let error_margin = f32::EPSILON;
    for e in entries.iter().rev() {
        let fee_rate = e.fee_per_byte();
        if bin_size > VSIZE_BIN_WIDTH && (last_fee_rate - fee_rate).abs() > error_margin {
            // vsize of transactions paying >= e.fee_per_vbyte()
            histogram.push((last_fee_rate, bin_size));
            bin_size = 0;
        }
        last_fee_rate = fee_rate;
        bin_size += e.size() as u32;
    }
    if bin_size > 0 {
        histogram.push((last_fee_rate, bin_size));
    }
    histogram
}

#[cfg(test)]
mod tests {
    use bitcoincash::blockdata::script::Builder;

    use crate::chaindef::{Transaction, TxOut};

    #[cfg(nexa)]
    fn one_output_tx() -> Transaction {
        use crate::nexa::transaction::TxOutType;
        let out = TxOut {
            txout_type: TxOutType::SATOSHI as u8,
            value: 0,
            script_pubkey: vec![].into(),
        };
        Transaction {
            version: 2,
            lock_time: 0,
            input: vec![],
            output: vec![out],
        }
    }

    #[cfg(bch)]
    fn one_output_tx() -> Transaction {
        let out = TxOut {
            value: 0,
            script_pubkey: vec![].into(),
            token: None,
        };
        Transaction {
            version: 2,
            lock_time: bitcoincash::PackedLockTime(0),
            input: vec![],
            output: vec![out],
        }
    }

    fn dummy_tx(desired_size: usize) -> Transaction {
        let mut tx = one_output_tx();
        let size = tx.size();
        let mut extra_bytes_size = desired_size.saturating_sub(size);
        let mut extra_bytes: Vec<u8> = vec![0; extra_bytes_size];
        tx.output[0].script_pubkey = Builder::from(extra_bytes).into_script();

        // tx uses varint for script_pubkey, so we need to reduce
        while tx.size() > desired_size {
            extra_bytes_size -= 1;
            extra_bytes = vec![0; extra_bytes_size];
            tx.output[0].script_pubkey = Builder::from(extra_bytes).into_script();
        }
        assert_eq!(tx.size(), desired_size);
        tx
    }

    #[test]
    fn test_fakestore() {
        use super::MempoolEntry;
        use crate::mempool::electrum_fees;

        let entries = [
            // (fee: u64, vsize: u32)
            &MempoolEntry::new(1_000, dummy_tx(1_000)),
            &MempoolEntry::new(10_000, dummy_tx(10_000)),
            &MempoolEntry::new(50_000, dummy_tx(50_000)),
            &MempoolEntry::new(120_000, dummy_tx(60_000)),
            &MempoolEntry::new(210_000, dummy_tx(70_000)),
            &MempoolEntry::new(320_000, dummy_tx(80_000)),
        ];
        assert_eq!(
            electrum_fees(&entries[..]),
            vec![(3.0, 150_000), (1.0, 121_000)]
        );
    }
}