rostrum 14.0.1

use crate::{
    chaindef::{Block, BlockHash, BlockHeader, OutPointHash, ScriptHash, Transaction},
    daemon::Daemon,
    encode::compute_outpoint_hash,
    index::index_block,
    index::last_indexed_block,
    indexes::{
        headerindex::HeaderRow, outputindex::OutputIndexRow,
        outtoscriptindex::OutToScripthashIndex, unspentindex::UnspentIndexRow, DBRow,
    },
    store::{DBStore, Row, META_CF},
    undo::blockundoer::BlockUndoer,
    writebatch::{PrefilledSpentData, SpentUpdateBatch, WriteBatch},
};
use anyhow::{Context, Result};
use async_trait::async_trait;
use bitcoin_hashes::hex::ToHex;
use bitcoincash::Txid;
use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
use std::{
    collections::{HashMap, HashSet},
    sync::{Arc, Mutex},
};

pub struct StoreBlockUndoer {
    store: Arc<DBStore>,
    daemon: Arc<Daemon>,
    header_chain_cache: Arc<Mutex<Option<Vec<BlockHeader>>>>,
}

impl StoreBlockUndoer {
    pub fn new(store: Arc<DBStore>, daemon: Arc<Daemon>) -> Self {
        Self {
            store,
            daemon,
            header_chain_cache: Arc::new(Mutex::new(None)),
        }
    }

    /// Get or build header chain for the given tip, using cache if available
    /// Invalidates cache if blockhash is not in the chain
    async fn get_header_chain(
        &self,
        tip: &BlockHash,
        blockhash: &BlockHash,
    ) -> Result<Vec<BlockHeader>> {
        // Check cache first
        {
            let mut cache = self.header_chain_cache.lock().unwrap();
            if let Some(ref cached_chain) = *cache {
                // Check if blockhash is in cached chain
                if cached_chain
                    .iter()
                    .any(|header| header.block_hash() == *blockhash)
                {
                    return Ok(cached_chain.clone());
                }
                // Blockhash not in chain - invalidate cache
                *cache = None;
            }
        }

        // Cache miss or invalid - build new chain
        let header_chain = HeaderRow::headers_up_to_tip(self.store.as_ref(), tip)
            .await
            .context("failed to build header chain for undo")?;

        // Update cache
        {
            let mut cache = self.header_chain_cache.lock().unwrap();
            *cache = Some(header_chain.clone());
        }

        Ok(header_chain)
    }

    /// Collect all outpoints created in the block
    fn collect_created_outputs(block: &Block) -> HashSet<OutPointHash> {
        block
            .txdata
            .par_iter()
            .map(|tx| {
                #[cfg(bch)]
                let txid_or_idem = tx.txid();
                #[cfg(nexa)]
                let txid_or_idem = tx.txidem();

                tx.output
                    .par_iter()
                    .enumerate()
                    .map(|(n, _out)| compute_outpoint_hash(&txid_or_idem, n as u32))
                    .collect::<Vec<OutPointHash>>()
            })
            .flatten()
            .collect()
    }

    /// Collect all outpoints spent in the block
    fn collect_spent_outputs(block: &Block) -> HashSet<OutPointHash> {
        block
            .txdata
            .par_iter()
            .map(|tx| {
                #[cfg(nexa)]
                let ophs = tx
                    .input
                    .iter()
                    .map(|input| input.previous_output.hash)
                    .collect::<Vec<OutPointHash>>();
                #[cfg(bch)]
                let ophs = tx
                    .input
                    .par_iter()
                    .map(|input| {
                        compute_outpoint_hash(
                            &input.previous_output.txid,
                            input.previous_output.vout,
                        )
                    })
                    .collect::<Vec<OutPointHash>>();

                ophs
            })
            .flatten()
            .collect()
    }

    /// Get scripthash and has_token for multiple outpoints in batch
    /// Groups by blockhash to avoid fetching the same block multiple times
    /// Returns HashMap suitable for PrefilledSpentData
    async fn get_scripthash_from_outpoints(
        &self,
        outpoints: &HashSet<OutPointHash>,
        blockhash: &BlockHash,
    ) -> Result<HashMap<OutPointHash, PrefilledSpentData>> {
        // Get or build header chain up to blockhash (the block being undone) for looking up blockhashes by height
        // Cache is invalidated if blockhash is not in the chain
        let header_chain = self.get_header_chain(blockhash, blockhash).await?;
        // Batch lookup all OutputIndexRows
        let mut output_rows = HashMap::new();
        for outpoint in outpoints {
            let output_key = OutputIndexRow::filter_by_outpointhash(outpoint);
            let (_, output_value) = self.store.get(OutputIndexRow::CF, output_key).await;
            if let Some(value) = output_value {
                let row = OutputIndexRow::from_row(&Row {
                    key: OutputIndexRow::filter_by_outpointhash(outpoint).into_boxed_slice(),
                    value: value.into_boxed_slice(),
                });
                output_rows.insert(*outpoint, row);
            }
        }

        // Group by blockhash to fetch each block only once
        let mut blockhash_to_outpoints: HashMap<BlockHash, Vec<(OutPointHash, Txid, u32)>> =
            HashMap::new();
        for (outpoint, output_row) in &output_rows {
            let txid = output_row.txid();
            let vout = output_row.index();
            let height = output_row.height() as usize;

            if let Some(header) = header_chain.get(height) {
                let blockhash = header.block_hash();
                blockhash_to_outpoints
                    .entry(blockhash)
                    .or_default()
                    .push((*outpoint, txid, vout));
            }
        }

        // Fetch blocks (order doesn't matter)
        let mut result = HashMap::new();
        info!(
            "undo: Fetching undo data from {} blocks for {} outpoints",
            blockhash_to_outpoints.len(),
            outpoints.len()
        );
        for (blockhash, outpoints_for_block) in blockhash_to_outpoints {
            // Fetch block from daemon
            let block = match self.daemon.getblock(&blockhash).await {
                Ok(block) => block,
                Err(e) => {
                    warn!(
                        "could not get block {} from daemon: {}, skipping {} outpoints",
                        blockhash.to_hex(),
                        e,
                        outpoints_for_block.len()
                    );
                    continue;
                }
            };

            // Build transaction cache for this block (to avoid O(n) lookup for every outpoint)
            let txcache: HashMap<Txid, &Transaction> =
                block.txdata.iter().map(|tx| (tx.txid(), tx)).collect();

            // Extract scripthash for all outpoints from this block
            for (outpoint, txid, vout) in outpoints_for_block {
                // Find the transaction in the block using cache
                if let Some(tx) = txcache.get(&txid) {
                    if let Some(txout) = tx.output.get(vout as usize) {
                        result.insert(
                            outpoint,
                            PrefilledSpentData {
                                scripthash: ScriptHash::normalized_from_txout(txout),
                                has_token: txout.has_token(),
                            },
                        );
                    } else {
                        warn!(
                            "skipping output {}: transaction {} does not have output index {}",
                            outpoint.to_hex(),
                            txid.to_hex(),
                            vout
                        );
                    }
                } else {
                    warn!(
                        "skipping output {}: could not find transaction {} in block {}",
                        outpoint.to_hex(),
                        txid.to_hex(),
                        blockhash.to_hex()
                    );
                }
            }
        }

        Ok(result)
    }

    /// Recreate unspent index entries that were spent in this block.
    fn recreate_unspent_data(
        &self,
        outpoints: &HashSet<OutPointHash>,
        prefilled_spent_data: &HashMap<OutPointHash, PrefilledSpentData>,
    ) -> Vec<Row> {
        outpoints
            .iter()
            .filter_map(|outpoint| {
                prefilled_spent_data.get(outpoint).map(|data| {
                    UnspentIndexRow::new(&data.scripthash, outpoint, data.has_token).to_row()
                })
            })
            .collect()
    }

    /// Recreate OutToScripthashIndex entries for outputs that were spent in
    /// this block but created in earlier blocks. These entries are pruned when
    /// an output is spent, and must be restored during undo so that the output
    /// can be re-spent when a new block is indexed.
    fn recreate_out_to_scripthash_data(
        outpoints: &HashSet<OutPointHash>,
        prefilled_spent_data: &HashMap<OutPointHash, PrefilledSpentData>,
    ) -> Vec<Row> {
        outpoints
            .iter()
            .filter_map(|outpoint| {
                prefilled_spent_data.get(outpoint).map(|data| {
                    OutToScripthashIndex::new(data.scripthash, *outpoint, data.has_token).to_row()
                })
            })
            .collect()
    }
}

#[async_trait]
impl BlockUndoer for StoreBlockUndoer {
    async fn undo_block(
        &self,
        blockhash: &BlockHash,
        height: u64,
        new_tip: &BlockHash,
    ) -> Result<()> {
        info!(
            "undo: Undoing block {} (height {})",
            blockhash.to_hex(),
            height
        );
        let block = self.daemon.getblock(blockhash).await?;

        let write = WriteBatch::new();
        let mut spentupdatebatch = SpentUpdateBatch::new().set_erase_only();
        let recreate_unspent = WriteBatch::new();

        // Prefill scripthash data for spent outputs (need to look up from earlier blocks)
        let spent_outputs = Self::collect_spent_outputs(&block);
        let prefilled_spent_data = self
            .get_scripthash_from_outpoints(&spent_outputs, blockhash)
            .await?;

        spentupdatebatch = spentupdatebatch.add_prefilled_spent_data(prefilled_spent_data.clone());

        index_block(&write, &spentupdatebatch, &block, height as usize);

        // genesis output is unspendable
        if height != 0 {
            // Recreate unspent entries for outputs created in earlier blocks but spent in this block,
            // (ignoring outputs created in this block)
            let created = Self::collect_created_outputs(&block);
            let spent = spent_outputs;

            // Recreate outputs spent in this block; but not created in this block.
            let old_spent: HashSet<OutPointHash> = spent.difference(&created).copied().collect();
            let recreated_unspent_rows =
                self.recreate_unspent_data(&old_spent, &prefilled_spent_data);
            recreate_unspent.insert(UnspentIndexRow::CF, recreated_unspent_rows);

            // Also recreate OutToScripthashIndex entries. These are pruned when
            // an output is spent and must be restored so that re-spending the
            // output in a future block can locate the scripthash.
            let recreated_scripthash_rows =
                Self::recreate_out_to_scripthash_data(&old_spent, &prefilled_spent_data);
            recreate_unspent.insert(OutToScripthashIndex::CF, recreated_scripthash_rows);
        };

        // undo; so erase
        self.store.update_spends(&spentupdatebatch);
        self.store.erase_batch(&write);

        // re-create utxos in block
        self.store.write_batch(&recreate_unspent);

        // The above erases the tip flag. Set the new tip.
        let batch = WriteBatch::new();
        batch.insert(META_CF, rayon::iter::once(last_indexed_block(new_tip)));
        self.store.write_batch(&batch);
        Ok(())
    }
}