blvm-consensus 0.1.82

//! Chain reorganization functions from Orange Paper Section 10.3

use crate::block::connect_block;
use crate::error::Result;
use crate::segwit::Witness;
use crate::types::*;
#[cfg(feature = "spec-lock")]
use blvm_spec_lock::spec_locked;
use std::collections::HashMap;

/// Deprecated reorg entry point (empty witnesses only).
///
/// Builds empty witness stacks suitable for coinbase-only / legacy blocks. Blocks containing
/// native SegWit transactions require explicit witness data via
/// [`reorganize_chain_with_witnesses`].
#[deprecated(
    since = "0.1.33",
    note = "Use reorganize_chain_with_witnesses with explicit witness data for SegWit chains"
)]
#[cfg_attr(feature = "spec-lock", spec_locked("11.3"))]
pub fn reorganize_chain(
    new_chain: &[Block],
    current_chain: &[Block],
    current_utxo_set: UtxoSet,
    current_height: Natural,
    network: crate::types::Network,
) -> Result<ReorganizationResult> {
    use crate::segwit::is_segwit_transaction;
    use crate::transaction::is_coinbase;

    for block in new_chain {
        for tx in &block.transactions {
            if !is_coinbase(tx) && is_segwit_transaction(tx) {
                return Err(crate::error::ConsensusError::BlockValidation(
                    "reorganize_chain: SegWit transactions require reorganize_chain_with_witnesses"
                        .into(),
                ));
            }
        }
    }

    if current_height > i64::MAX as u64 {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!("Current height {current_height} must fit in i64").into(),
        ));
    }
    if current_utxo_set.len() > u32::MAX as usize {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!(
                "Current UTXO set size {} exceeds maximum",
                current_utxo_set.len()
            )
            .into(),
        ));
    }

    // Empty per-input witness stacks (legacy path only; SegWit rejected above).
    let empty_witnesses: Vec<Vec<Vec<Witness>>> = new_chain
        .iter()
        .map(|block| {
            block
                .transactions
                .iter()
                .map(|tx| tx.inputs.iter().map(|_| Vec::new()).collect())
                .collect()
        })
        .collect();

    reorganize_chain_with_witnesses(
        new_chain,
        &empty_witnesses,
        None, // No headers for median time-past
        current_chain,
        current_utxo_set,
        current_height,
        None::<fn(&Block) -> Option<Vec<Witness>>>, // No witness retrieval
        None::<fn(Natural) -> Option<Vec<BlockHeader>>>, // No header retrieval
        None::<fn(&Hash) -> Option<BlockUndoLog>>,  // No undo log retrieval
        None::<fn(&Hash, &BlockUndoLog) -> Result<()>>, // No undo log storage
        new_chain
            .iter()
            .map(|b| b.header.timestamp)
            .max()
            .unwrap_or(0)
            .saturating_add(crate::constants::MAX_FUTURE_BLOCK_TIME),
        network,
        None,
    )
}

/// Reorganization: When a longer chain is found (full API with witness support)
///
/// For new chain with blocks [b1, b2, ..., bn] and current chain with blocks [c1, c2, ..., cm]:
/// 1. Find common ancestor between new chain and current chain
/// 2. Disconnect blocks from current chain back to common ancestor
/// 3. Connect blocks from new chain from common ancestor forward
/// 4. Return new UTXO set and reorganization result
///
/// # Arguments
///
/// * `new_chain` - Blocks from the new (longer) chain
/// * `new_chain_witnesses` - Witness data for each block in new_chain (one Vec<Witness> per block)
/// * `new_chain_headers` - Recent headers for median time-past calculation (last 11+ headers, oldest to newest)
/// * `current_chain` - Blocks from the current chain
/// * `current_utxo_set` - Current UTXO set
/// * `current_height` - Current block height
/// * `get_witnesses_for_block` - Optional callback to retrieve witnesses for a block (for current chain disconnection)
/// * `get_headers_for_height` - Optional callback to retrieve headers for median time-past (for current chain disconnection)
/// * `get_undo_log_for_block` - Optional callback to retrieve undo log for a block (for current chain disconnection)
/// * `store_undo_log_for_block` - Optional callback to store undo log for a block (for new chain connection)
/// * `network_time` - Adjusted network time from the node layer (not the block under validation)
/// * `network` - Consensus network (mainnet / testnet / regtest)
#[allow(clippy::too_many_arguments)]
#[cfg_attr(feature = "spec-lock", spec_locked("11.3"))]
pub fn reorganize_chain_with_witnesses(
    new_chain: &[Block],
    new_chain_witnesses: &[Vec<Vec<Witness>>], // CRITICAL FIX: Changed from &[Vec<Witness>] to &[Vec<Vec<Witness>>]
    // Each block has Vec<Vec<Witness>> (one Vec per transaction, each containing one Witness per input)
    new_chain_headers: Option<&[BlockHeader]>,
    current_chain: &[Block],
    current_utxo_set: UtxoSet,
    current_height: Natural,
    _get_witnesses_for_block: Option<impl Fn(&Block) -> Option<Vec<Witness>>>,
    get_headers_for_height: Option<impl Fn(Natural) -> Option<Vec<BlockHeader>>>,
    get_undo_log_for_block: Option<impl Fn(&Hash) -> Option<BlockUndoLog>>,
    store_undo_log_for_block: Option<impl Fn(&Hash, &BlockUndoLog) -> Result<()>>,
    network_time: u64,
    network: crate::types::Network,
    mut connect_context_for_height: Option<
        &mut dyn FnMut(
            u64,
            Option<&[BlockHeader]>,
            u64,
            crate::types::Network,
        ) -> crate::block::BlockValidationContext,
    >,
) -> Result<ReorganizationResult> {
    if current_height > i64::MAX as u64 {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!("Current height {current_height} must fit in i64").into(),
        ));
    }
    if current_utxo_set.len() > u32::MAX as usize {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!(
                "Current UTXO set size {} exceeds maximum",
                current_utxo_set.len()
            )
            .into(),
        ));
    }
    if new_chain_witnesses.len() != new_chain.len() {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!(
                "New chain witness count {} must match block count {}",
                new_chain_witnesses.len(),
                new_chain.len()
            )
            .into(),
        ));
    }

    // 1. Find common ancestor by comparing block hashes
    let common_ancestor = find_common_ancestor(new_chain, current_chain)?;
    let common_ancestor_header = common_ancestor.header;
    let common_ancestor_index = common_ancestor.new_chain_index;
    let current_ancestor_index = common_ancestor.current_chain_index;

    if common_ancestor_index >= new_chain.len() {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!(
                "Common ancestor index {common_ancestor_index} must be < new chain length {}",
                new_chain.len()
            )
            .into(),
        ));
    }
    if current_ancestor_index >= current_chain.len() {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!(
                "Common ancestor index {current_ancestor_index} must be < current chain length {}",
                current_chain.len()
            )
            .into(),
        ));
    }

    // 2. Disconnect blocks from current chain back to common ancestor
    // We disconnect from (current_ancestor_index + 1) to the tip
    // Undo logs are retrieved from persistent storage via the callback.
    // The node layer (blvm-node) should provide a callback that uses BlockStore::get_undo_log()
    // to retrieve undo logs from the database (redb/sled).
    let mut utxo_set = current_utxo_set;
    if utxo_set.len() > u32::MAX as usize {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!("UTXO set size {} must not exceed maximum", utxo_set.len()).into(),
        ));
    }

    // Disconnect from the block after the common ancestor to the tip
    let disconnect_start = current_ancestor_index + 1;
    if disconnect_start > current_chain.len() {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!(
                "Disconnect start {disconnect_start} must be <= current chain length {}",
                current_chain.len()
            )
            .into(),
        ));
    }

    let mut disconnected_undo_logs: HashMap<Hash, BlockUndoLog> = HashMap::new();

    for i in (disconnect_start..current_chain.len()).rev() {
        if let Some(block) = current_chain.get(i) {
            if block.transactions.is_empty() {
                return Err(crate::error::ConsensusError::BlockValidation(
                    format!("Block at index {i} must have at least one transaction").into(),
                ));
            }

            let block_hash = calculate_block_hash(&block.header);
            if block_hash == [0u8; 32] {
                return Err(crate::error::ConsensusError::BlockValidation(
                    "Block hash must be non-zero".into(),
                ));
            }

            // The node stores undo logs under the 80-byte block id. A miss must
            // not disconnect with an empty log: that leaves the old fork's coins
            // spendable and the coins it spent still spent.
            let undo_log = if let Some(ref get_undo_log) = get_undo_log_for_block {
                get_undo_log(&block_hash).ok_or_else(|| {
                    crate::error::ConsensusError::BlockValidation(
                        format!("missing undo log for block {block_hash:?}").into(),
                    )
                })?
            } else {
                return Err(crate::error::ConsensusError::BlockValidation(
                    "missing undo log for disconnected block".into(),
                ));
            };

            utxo_set = disconnect_block(block, &undo_log, utxo_set, (i as Natural) + 1)?;
            disconnected_undo_logs.insert(block_hash, undo_log);
        }
    }

    // 3. Connect blocks from new chain from common ancestor forward
    // We connect from (common_ancestor_index + 1) to the tip of new chain
    // Calculate the height at the common ancestor.
    // current_chain[i] is at height: current_height - (current_chain.len() - 1 - i)
    // So ancestor at current_ancestor_index is at:
    //   current_height - (current_chain.len() - 1 - current_ancestor_index)
    let blocks_after_ancestor = (current_chain.len() - 1 - current_ancestor_index) as Natural;
    let common_ancestor_height = current_height.saturating_sub(blocks_after_ancestor);
    let mut new_height = common_ancestor_height;
    let mut connected_blocks = Vec::new();
    let mut connected_undo_logs: HashMap<Hash, BlockUndoLog> = HashMap::new();
    let mut mtp_header_buf: Vec<BlockHeader> = Vec::new();

    // Ensure witnesses match blocks
    if new_chain_witnesses.len() != new_chain.len() {
        return Err(crate::error::ConsensusError::ConsensusRuleViolation(
            format!(
                "Witness count {} does not match block count {}",
                new_chain_witnesses.len(),
                new_chain.len()
            )
            .into(),
        ));
    }

    // Connect blocks starting from the block after the common ancestor
    for (i, block) in new_chain.iter().enumerate().skip(common_ancestor_index + 1) {
        new_height += 1;
        // Witness stacks are required; length is validated at entry (must match new_chain).
        let witnesses = new_chain_witnesses[i].clone();

        // Median time-past ends at height-1. A supplied parent window is the
        // previous 11 headers. The in-fork slice is only a fallback for callers
        // that have no header store; a one-block fork does not contain those parents.
        mtp_header_buf.clear();
        let expected = if new_height == 0 {
            0
        } else {
            (new_height as usize).min(crate::bip113::MEDIAN_TIME_BLOCKS)
        };
        let stored_headers = get_headers_for_height
            .as_ref()
            .and_then(|get| get(new_height));
        let recent_headers: Option<&[BlockHeader]> = if let Some(headers) = stored_headers.as_ref()
        {
            if headers.len() < expected {
                return Err(crate::error::ConsensusError::BlockValidation(
                    format!("missing parent headers for median time at height {new_height}").into(),
                ));
            }
            mtp_header_buf.extend(headers.iter().cloned());
            Some(mtp_header_buf.as_slice())
        } else if let Some(headers) = new_chain_headers {
            Some(headers)
        } else if i > 0 {
            let start = i.saturating_sub(crate::bip113::MEDIAN_TIME_BLOCKS - 1);
            mtp_header_buf.extend(new_chain[start..i].iter().map(|b| b.header.clone()));
            Some(mtp_header_buf.as_slice())
        } else if new_height == 0 {
            Some(&[])
        } else {
            None
        };

        let context = if let Some(build) = connect_context_for_height.as_mut() {
            build(new_height, recent_headers, network_time, network)
        } else {
            crate::block::block_validation_context_for_connect_ibd(
                recent_headers,
                network_time,
                network,
            )
        };
        let (validation_result, new_utxo_set, undo_log) =
            connect_block(block, &witnesses, utxo_set, new_height, &context)?;

        if let ValidationResult::Invalid(reason) = validation_result {
            return Err(crate::error::ConsensusError::ConsensusRuleViolation(
                format!("Invalid block at height {new_height} during reorganization: {reason}")
                    .into(),
            ));
        }

        // Store undo log for this block (keyed by block hash for future retrieval)
        let block_hash = calculate_block_hash(&block.header);

        // Persist undo log to database via callback (required for future reorganizations)
        if let Some(ref store_undo_log) = store_undo_log_for_block {
            store_undo_log(&block_hash, &undo_log)?;
        }

        // Also store in-memory for the reorganization result
        connected_undo_logs.insert(block_hash, undo_log);

        utxo_set = new_utxo_set;
        connected_blocks.push(block.clone());
    }

    // 4. Return reorganization result
    Ok(ReorganizationResult {
        new_utxo_set: utxo_set,
        new_height,
        common_ancestor: common_ancestor_header,
        disconnected_blocks: current_chain[disconnect_start..].to_vec(),
        connected_blocks,
        reorganization_depth: current_chain.len() - disconnect_start,
        connected_block_undo_logs: connected_undo_logs,
    })
}

/// Update mempool after chain reorganization
///
/// This function should be called after successfully reorganizing the chain
/// to keep the mempool synchronized with the new blockchain state.
///
/// Handles:
/// 1. Removes transactions from mempool that were included in the new chain blocks
/// 2. Removes transactions that became invalid (their inputs were spent by new chain)
/// 3. Optionally re-adds transactions from disconnected blocks that are still valid
///
/// # Arguments
///
/// * `mempool` - Mutable reference to the mempool
/// * `reorg_result` - The reorganization result
/// * `utxo_set` - The updated UTXO set after reorganization
/// * `get_tx_by_id` - Optional function to look up transactions by ID (needed for full validation)
///
/// # Returns
///
/// Returns a vector of transaction IDs that were removed from the mempool.
///
/// # Example
///
/// ```rust
/// use blvm_consensus::reorganization::{reorganize_chain_with_witnesses, update_mempool_after_reorg};
/// use blvm_consensus::mempool::Mempool;
/// use blvm_consensus::segwit::Witness;
///
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// # use blvm_consensus::types::*;
/// # use blvm_consensus::mempool::Mempool;
/// # let new_chain = vec![];
/// # let new_witnesses = vec![];
/// # let current_chain = vec![];
/// # let current_utxo_set = UtxoSet::default();
/// # let current_height = 0;
/// # let mut mempool = Mempool::new();
/// // Note: This is a simplified example. In practice, chains must have at least one block
/// // and share a common ancestor. The result may be an error for empty chains.
/// let reorg_result = reorganize_chain_with_witnesses(
///     &new_chain,
///     &new_witnesses,
///     None,
///     &current_chain,
///     current_utxo_set,
///     current_height,
///     None::<fn(&Block) -> Option<Vec<Witness>>>,
///     None::<fn(Natural) -> Option<Vec<BlockHeader>>>,
///     None::<fn(&blvm_consensus::types::Hash) -> Option<blvm_consensus::reorganization::BlockUndoLog>>,
///     None::<fn(&blvm_consensus::types::Hash, &blvm_consensus::reorganization::BlockUndoLog) -> blvm_consensus::error::Result<()>>,
///     0,
///     Network::Regtest,
///     None,
/// );
/// if let Ok(reorg_result) = reorg_result {
///     let _removed = update_mempool_after_reorg(
///         &mut mempool,
///         &reorg_result,
///         &reorg_result.new_utxo_set,
///         None::<fn(&Hash) -> Option<Transaction>>, // No transaction lookup available
///     )?;
/// }
/// # Ok(())
/// # }
/// ```
pub fn update_mempool_after_reorg<F>(
    mempool: &mut crate::mempool::Mempool,
    reorg_result: &ReorganizationResult,
    utxo_set: &UtxoSet,
    get_tx_by_id: Option<F>,
) -> Result<Vec<Hash>>
where
    F: Fn(&Hash) -> Option<Transaction>,
{
    use crate::mempool::update_mempool_after_block;

    let mut all_removed = Vec::new();

    // 1. Remove transactions that were included in the new connected blocks
    for block in &reorg_result.connected_blocks {
        let removed = update_mempool_after_block(mempool, block, utxo_set)?;
        all_removed.extend(removed);
    }

    // 2. Remove transactions that became invalid (their inputs were spent by new chain)
    // Collect all spent outpoints from the new connected blocks
    let mut spent_outpoints = std::collections::HashSet::new();
    for block in &reorg_result.connected_blocks {
        for tx in &block.transactions {
            if !crate::transaction::is_coinbase(tx) {
                for input in &tx.inputs {
                    spent_outpoints.insert(input.prevout);
                }
            }
        }
    }

    // If we have transaction lookup, check each mempool transaction
    if let Some(lookup) = get_tx_by_id {
        let mut invalid_tx_ids = Vec::new();
        for &tx_id in mempool.iter() {
            if let Some(tx) = lookup(&tx_id) {
                // Check if any input of this transaction was spent by the new chain
                for input in &tx.inputs {
                    if spent_outpoints.contains(&input.prevout) {
                        invalid_tx_ids.push(tx_id);
                        break;
                    }
                }
            }
        }

        // Remove invalid transactions
        for tx_id in invalid_tx_ids {
            if mempool.remove(&tx_id) {
                all_removed.push(tx_id);
            }
        }
    }

    // 3. Re-add transactions from disconnected blocks via [`update_mempool_after_reorg_with_readd`].

    Ok(all_removed)
}

/// Parameters for re-adding mempool entries from disconnected blocks after a reorg.
#[derive(Debug, Clone)]
pub struct ReorgMempoolReaddParams {
    pub height: Natural,
    pub time_context: Option<TimeContext>,
    pub network: Network,
}

/// Update mempool after reorg and re-add valid transactions from disconnected blocks.
///
/// Runs [`update_mempool_after_reorg`] then attempts `accept_to_memory_pool` for each
/// non-coinbase transaction in `reorg_result.disconnected_blocks`.
pub fn update_mempool_after_reorg_with_readd<F, G>(
    mempool: &mut crate::mempool::Mempool,
    reorg_result: &ReorganizationResult,
    utxo_set: &UtxoSet,
    get_tx_by_id: Option<F>,
    readd: &ReorgMempoolReaddParams,
    get_witnesses: Option<G>,
) -> Result<(Vec<Hash>, Vec<Hash>)>
where
    F: Fn(&Hash) -> Option<Transaction>,
    G: Fn(&Hash) -> Option<Vec<Witness>>,
{
    let removed = update_mempool_after_reorg(mempool, reorg_result, utxo_set, get_tx_by_id)?;
    let readded = readd_disconnected_to_mempool(
        mempool,
        &reorg_result.disconnected_blocks,
        utxo_set,
        readd,
        get_witnesses,
    )?;
    Ok((removed, readded))
}

fn readd_disconnected_to_mempool<G>(
    mempool: &mut crate::mempool::Mempool,
    disconnected_blocks: &[Block],
    utxo_set: &UtxoSet,
    readd: &ReorgMempoolReaddParams,
    get_witnesses: Option<G>,
) -> Result<Vec<Hash>>
where
    G: Fn(&Hash) -> Option<Vec<Witness>>,
{
    use crate::block::calculate_tx_id;
    use crate::mempool::{MempoolResult, accept_to_memory_pool};
    use crate::transaction::is_coinbase;

    let mut readded = Vec::new();
    for block in disconnected_blocks {
        for tx in block.transactions.iter() {
            if is_coinbase(tx) {
                continue;
            }
            let tx_id = calculate_tx_id(tx);
            if mempool.contains(&tx_id) {
                continue;
            }
            let witness_storage = get_witnesses.as_ref().and_then(|lookup| lookup(&tx_id));
            let witnesses = witness_storage.as_deref();
            match accept_to_memory_pool(
                tx,
                witnesses,
                utxo_set,
                mempool,
                readd.height,
                readd.time_context,
                readd.network,
                None,
            )? {
                MempoolResult::Accepted => {
                    if mempool.insert_transaction_with_witness(tx, witnesses) {
                        readded.push(tx_id);
                    }
                }
                MempoolResult::Rejected(_) => {}
            }
        }
    }
    Ok(readded)
}

/// Update mempool after reorg without a transaction lookup callback.
pub fn update_mempool_after_reorg_simple(
    mempool: &mut crate::mempool::Mempool,
    reorg_result: &ReorganizationResult,
    utxo_set: &UtxoSet,
) -> Result<Vec<Hash>> {
    update_mempool_after_reorg(
        mempool,
        reorg_result,
        utxo_set,
        None::<fn(&Hash) -> Option<Transaction>>,
    )
}

/// Common ancestor result with indices in both chains
struct CommonAncestorResult {
    header: BlockHeader,
    new_chain_index: usize,
    current_chain_index: usize,
}

/// Find common ancestor between two chains by comparing block hashes.
///
/// Both slices are oldest-to-tip and may start at different heights. The fork
/// point is the shared block nearest the new tip, not the block at the same
/// index. Index comparison rejects a valid heavier chain once either tip is
/// far enough from genesis that the windows no longer start on the same block.
/// Orange Paper 11.3: Chain reorganization finds common ancestor before disconnect/connect.
/// Height of the common ancestor when reorganizing from `current_chain` at `current_height`.
pub fn common_ancestor_height_at_reorg(
    current_height: u64,
    current_chain: &[Block],
    new_chain: &[Block],
) -> Result<u64> {
    let ca = find_common_ancestor(new_chain, current_chain)?;
    let blocks_after = current_chain
        .len()
        .saturating_sub(1)
        .saturating_sub(ca.current_chain_index);
    Ok(current_height.saturating_sub(blocks_after as u64))
}

fn find_common_ancestor(
    new_chain: &[Block],
    current_chain: &[Block],
) -> Result<CommonAncestorResult> {
    if new_chain.is_empty() || current_chain.is_empty() {
        return Err(crate::error::ConsensusError::ConsensusRuleViolation(
            "Cannot find common ancestor: empty chain".into(),
        ));
    }

    let mut current_at: HashMap<Hash, usize> = HashMap::with_capacity(current_chain.len());
    for (i, block) in current_chain.iter().enumerate() {
        current_at.insert(calculate_block_hash(&block.header), i);
    }

    let mut found: Option<(usize, usize)> = None;
    for (new_i, block) in new_chain.iter().enumerate() {
        let hash = calculate_block_hash(&block.header);
        if let Some(&cur_i) = current_at.get(&hash) {
            found = Some((new_i, cur_i));
        }
    }

    let Some((new_i, cur_i)) = found else {
        return Err(crate::error::ConsensusError::ConsensusRuleViolation(
            "Chains do not share a common ancestor".into(),
        ));
    };
    Ok(CommonAncestorResult {
        header: new_chain[new_i].header.clone(),
        new_chain_index: new_i,
        current_chain_index: cur_i,
    })
}

/// Disconnect a block from the chain (reverse of ConnectBlock)
///
/// Uses the undo log to perfectly restore the UTXO set to its state before the block was connected.
/// This is the inverse operation of `connect_block`.
/// Orange Paper 11.3.1: DisconnectBlock applies undo log to reverse ConnectBlock.
///
/// # Arguments
///
/// * `block` - The block to disconnect (used for validation, undo_log contains the actual changes)
/// * `undo_log` - The undo log created when this block was connected
/// * `utxo_set` - Current UTXO set (will be modified)
/// * `_height` - Block height (for potential future use)
#[cfg_attr(feature = "spec-lock", spec_locked("11.3.1", "DisconnectBlock"))]
fn disconnect_block(
    _block: &Block,
    undo_log: &BlockUndoLog,
    mut utxo_set: UtxoSet,
    _height: Natural,
) -> Result<UtxoSet> {
    if _block.transactions.is_empty() {
        return Err(crate::error::ConsensusError::BlockValidation(
            "Block must have at least one transaction".into(),
        ));
    }
    if _height > i64::MAX as u64 {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!("Block height {_height} must fit in i64").into(),
        ));
    }
    if utxo_set.len() > u32::MAX as usize {
        return Err(crate::error::ConsensusError::BlockValidation(
            format!("UTXO set size {} must not exceed maximum", utxo_set.len()).into(),
        ));
    }
    // Process undo entries in reverse order (most recent first)
    // This reverses the order of operations from connect_block
    for entry in undo_log.entries.iter() {
        // Remove new UTXO (if it was created by this block)
        if entry.new_utxo.is_some() {
            utxo_set.remove(&entry.outpoint);
        }

        // Restore previous UTXO (if it was spent by this block)
        if let Some(previous_utxo) = &entry.previous_utxo {
            utxo_set.insert(entry.outpoint, std::sync::Arc::clone(previous_utxo));
        }
    }

    Ok(utxo_set)
}

/// Undo log a connect would have stored, rebuilt from the block and its spent outputs.
///
/// Entry order matches `connect_block`: spends then new outputs, then reversed so
/// disconnect walks the most recent change first. A missing prevout is an error.
/// An empty log is not a substitute for a block that spent or created coins.
pub fn undo_log_from_connected_block(
    block: &Block,
    height: Natural,
    mut lookup_prevout: impl FnMut(&OutPoint) -> Result<UTXO>,
) -> Result<BlockUndoLog> {
    use crate::block::calculate_tx_id;
    use crate::transaction::is_coinbase;

    if block.transactions.is_empty() {
        return Err(crate::error::ConsensusError::BlockValidation(
            "Block must have at least one transaction".into(),
        ));
    }
    let mut entries = Vec::new();
    for tx in block.transactions.iter() {
        let tx_id = calculate_tx_id(tx);
        if !is_coinbase(tx) {
            for input in tx.inputs.iter() {
                let previous = lookup_prevout(&input.prevout)?;
                entries.push(UndoEntry {
                    outpoint: input.prevout,
                    previous_utxo: Some(std::sync::Arc::new(previous)),
                    new_utxo: None,
                });
            }
        }
        for (index, output) in tx.outputs.iter().enumerate() {
            let created = UTXO {
                value: output.value,
                script_pubkey: output.script_pubkey.as_slice().into(),
                height,
                is_coinbase: is_coinbase(tx),
            };
            entries.push(UndoEntry {
                outpoint: OutPoint {
                    hash: tx_id,
                    index: index as u32,
                },
                previous_utxo: None,
                new_utxo: Some(std::sync::Arc::new(created)),
            });
        }
    }
    entries.reverse();
    Ok(BlockUndoLog { entries })
}

/// Check if reorganization is beneficial
#[track_caller] // Better error messages showing caller location
#[allow(clippy::redundant_comparisons)] // Intentional assertions for formal verification
#[cfg_attr(feature = "spec-lock", spec_locked("11.3", "ShouldReorganize"))]
pub fn should_reorganize(new_chain: &[Block], current_chain: &[Block]) -> Result<bool> {
    // Reorganize when new chain has strictly more cumulative work.
    let new_work = calculate_chain_work(new_chain)?;
    let current_work = calculate_chain_work(current_chain)?;
    Ok(new_work > current_work)
}

/// Work contribution for a single block header (matches [`calculate_chain_work`]).
pub fn work_for_bits(bits: Natural) -> Result<crate::pow::U256> {
    crate::pow::get_block_proof(bits)
}

/// Calculate total work for a chain
/// Orange Paper 11.3: BestChain = argmax Σ Work(block)
///
/// Mathematical invariants:
/// - Work is always non-negative
/// - Work increases monotonically with chain length
/// - Work calculation is deterministic
#[cfg_attr(feature = "spec-lock", spec_locked("11.3", "CalculateChainWork"))]
pub fn calculate_chain_work(chain: &[Block]) -> Result<crate::pow::U256> {
    let mut total_work = crate::pow::U256::zero();

    for block in chain {
        let work_contribution = work_for_bits(block.header.bits)?;
        let old_total = total_work;
        total_work = total_work.saturating_add(work_contribution);

        // Runtime assertion: Total work must be non-decreasing
        debug_assert!(
            total_work >= old_total,
            "Total work ({total_work:?}) must be >= previous total ({old_total:?})"
        );
    }

    Ok(total_work)
}

/// Test-only approximate tx id from tx fields (not a consensus hash).
#[allow(dead_code)] // Used in tests
fn calculate_tx_id(tx: &Transaction) -> Hash {
    let mut hash = [0u8; 32];
    hash[0] = (tx.version & 0xff) as u8;
    hash[1] = (tx.inputs.len() & 0xff) as u8;
    hash[2] = (tx.outputs.len() & 0xff) as u8;
    hash[3] = (tx.lock_time & 0xff) as u8;
    hash
}

/// Undo-log key. This is the 80-byte header hash, the same id the node stores.
fn calculate_block_hash(header: &BlockHeader) -> Hash {
    crate::block::block_header_hash(header)
}

// ============================================================================
// TYPES
// ============================================================================

mod undo_entry_serde {
    use crate::types::UTXO;
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use std::sync::Arc;

    pub fn serialize<S>(opt: &Option<Arc<UTXO>>, s: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        opt.as_ref().map(|a| a.as_ref()).serialize(s)
    }

    pub fn deserialize<'de, D>(d: D) -> Result<Option<Arc<UTXO>>, D::Error>
    where
        D: Deserializer<'de>,
    {
        Option::<UTXO>::deserialize(d).map(|opt| opt.map(Arc::new))
    }
}

/// Undo log entry for a single UTXO change
///
/// Records the state of a UTXO before and after a transaction is applied.
/// This allows perfect reversal of UTXO set changes during block disconnection.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct UndoEntry {
    /// The outpoint that was changed
    pub outpoint: OutPoint,
    /// The UTXO that existed before (None if it was created by this transaction)
    #[serde(with = "undo_entry_serde")]
    pub previous_utxo: Option<std::sync::Arc<UTXO>>,
    /// The UTXO that exists after (None if it was spent by this transaction)
    #[serde(with = "undo_entry_serde")]
    pub new_utxo: Option<std::sync::Arc<UTXO>>,
}

/// Undo log for a single block
///
/// Contains all UTXO changes made by a block, allowing perfect reversal
/// of the block's effects on the UTXO set.
///
/// Entries are stored in reverse order (most recent first) to allow
/// efficient undo by iterating forward.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct BlockUndoLog {
    /// Entries in reverse order (most recent first)
    /// This allows efficient undo: iterate forward and restore previous_utxo, remove new_utxo
    pub entries: Vec<UndoEntry>,
}

impl BlockUndoLog {
    /// Create an empty undo log
    pub fn new() -> Self {
        Self {
            entries: Vec::new(),
        }
    }

    /// Add an undo entry to the log
    pub fn push(&mut self, entry: UndoEntry) {
        self.entries.push(entry);
    }

    /// Check if the undo log is empty
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }
}

impl Default for BlockUndoLog {
    fn default() -> Self {
        Self::new()
    }
}

/// Result of chain reorganization
#[derive(Debug, Clone)]
pub struct ReorganizationResult {
    pub new_utxo_set: UtxoSet,
    pub new_height: Natural,
    pub common_ancestor: BlockHeader,
    pub disconnected_blocks: Vec<Block>,
    pub connected_blocks: Vec<Block>,
    pub reorganization_depth: usize,
    /// Undo logs for connected blocks (keyed by block hash)
    /// These can be used for future disconnections
    pub connected_block_undo_logs: HashMap<Hash, BlockUndoLog>,
}

// ============================================================================
// FORMAL VERIFICATION
// ============================================================================

/// Mathematical Specification for Chain Selection:
/// ∀ chains C₁, C₂: work(C₁) > work(C₂) ⇒ select(C₁)
///
/// Invariants:
/// - Selected chain has maximum cumulative work
/// - Work calculation is deterministic
/// - Empty chains are rejected
/// - Chain work is always non-negative

#[cfg(test)]
mod property_tests {
    use super::*;
    use proptest::prelude::*;

    /// Helper to get chain length range based on coverage mode
    fn chain_len_range() -> std::ops::Range<usize> {
        if std::env::var("CARGO_TARPAULIN").is_ok() || std::env::var("TARPAULIN").is_ok() {
            1..3 // Reduced range under coverage
        } else {
            1..5
        }
    }

    /// Helper to get chain length range for deterministic test
    fn chain_len_range_det() -> std::ops::Range<usize> {
        if std::env::var("CARGO_TARPAULIN").is_ok() || std::env::var("TARPAULIN").is_ok() {
            0..3 // Reduced range under coverage
        } else {
            0..10
        }
    }

    fn compact_target_bytes(exponent: u8, mantissa: u32) -> [u8; 32] {
        let mut out = [0u8; 32];
        if mantissa == 0 || !(3..=32).contains(&exponent) {
            return out;
        }
        let shift = (exponent as usize).saturating_sub(3);
        for (i, byte) in mantissa.to_le_bytes().iter().enumerate() {
            let pos = shift + i;
            if pos < 32 {
                out[pos] = *byte;
            }
        }
        out
    }

    /// Header bits whose compact exponent is 3..=32, so `expand_target` succeeds.
    fn arb_block() -> impl Strategy<Value = Block> {
        (3u8..=32, 0u32..=0x007f_ffff).prop_map(|(exponent, mantissa)| Block {
            header: BlockHeader {
                version: 1,
                prev_block_hash: [0u8; 32],
                merkle_root: [0u8; 32],
                timestamp: 0,
                bits: ((exponent as u64) << 24) | (mantissa as u64),
                nonce: 0,
            },
            transactions: Box::new([]),
        })
    }

    /// Property test: should_reorganize selects chain with maximum work
    proptest! {
        #[test]
        fn prop_should_reorganize_max_work(
            new_chain in proptest::collection::vec(arb_block(), chain_len_range()),
            current_chain in proptest::collection::vec(arb_block(), chain_len_range())
        ) {
            let new_w = calculate_chain_work(&new_chain).unwrap();
            let current_w = calculate_chain_work(&current_chain).unwrap();
            let should_reorg = should_reorganize(&new_chain, &current_chain).unwrap();
            prop_assert_eq!(should_reorg, new_w > current_w);
        }
    }

    /// Property test: calculate_chain_work is deterministic
    proptest! {
        #[test]
        fn prop_calculate_chain_work_deterministic(
            chain in proptest::collection::vec(arb_block(), chain_len_range_det())
        ) {
            let work1 = calculate_chain_work(&chain);
            let work2 = calculate_chain_work(&chain);
            prop_assert_eq!(work1, work2);
        }
    }

    /// Property test: expand_target handles various difficulty values
    proptest! {
        #[test]
        fn prop_expand_target_valid_range(
            exponent in 0u8..=40,
            mantissa in 0u32..=0x007f_ffff
        ) {
            let bits = ((exponent as u64) << 24) | (mantissa as u64);
            let result = crate::pow::expand_target(bits);
            if !(3..=32).contains(&exponent) {
                prop_assert!(result.is_err());
            } else {
                let expected = compact_target_bytes(exponent, mantissa);
                prop_assert_eq!(result.unwrap().to_le_bytes(), expected);
                let signed = bits | 0x0080_0000;
                prop_assert_eq!(
                    crate::pow::expand_target(signed).unwrap().to_le_bytes(),
                    expected
                );
            }
        }
    }

    /// Property test: should_reorganize with equal length chains compares work
    proptest! {
        #[test]
        fn prop_should_reorganize_equal_length(
            chain1 in proptest::collection::vec(arb_block(), 1..3),
            chain2 in proptest::collection::vec(arb_block(), 1..3)
        ) {
            // Ensure equal length
            let len = chain1.len().min(chain2.len());
            let chain1 = &chain1[..len];
            let chain2 = &chain2[..len];

            let w1 = calculate_chain_work(chain1).unwrap();
            let w2 = calculate_chain_work(chain2).unwrap();
            let should_reorg = should_reorganize(chain1, chain2).unwrap();
            prop_assert_eq!(should_reorg, w1 > w2);
        }
    }
}

#[cfg(test)]
#[path = "reorganization_tests.rs"]
mod tests;

#[cfg(test)]
mod undo_rebuild_tests {
    use super::*;
    use crate::block::calculate_tx_id;
    use crate::transaction::is_coinbase;
    use std::sync::Arc;

    fn coinbase_with_value(value: i64) -> Transaction {
        Transaction {
            version: 1,
            inputs: vec![TransactionInput {
                prevout: OutPoint {
                    hash: [0; 32],
                    index: 0xffffffff,
                },
                script_sig: vec![0x04].into(),
                sequence: 0xffffffff,
            }]
            .into(),
            outputs: vec![TransactionOutput {
                value,
                script_pubkey: vec![0x51],
            }]
            .into(),
            lock_time: 0,
        }
    }

    #[test]
    fn rebuilt_undo_drops_created_coins_and_restores_the_spend() {
        let parent_tx = coinbase_with_value(50);
        let parent_id = calculate_tx_id(&parent_tx);
        let spent = OutPoint {
            hash: parent_id,
            index: 0,
        };
        let previous = UTXO {
            value: 50,
            script_pubkey: vec![0x51].into(),
            height: 4,
            is_coinbase: true,
        };
        let spender = Transaction {
            version: 1,
            inputs: vec![TransactionInput {
                prevout: spent,
                script_sig: vec![].into(),
                sequence: 0xffffffff,
            }]
            .into(),
            outputs: vec![TransactionOutput {
                value: 40,
                script_pubkey: vec![0x52],
            }]
            .into(),
            lock_time: 0,
        };
        let block = Block {
            header: BlockHeader {
                version: 1,
                prev_block_hash: [1; 32],
                merkle_root: [2; 32],
                timestamp: 1,
                bits: 0x207fffff,
                nonce: 0,
            },
            transactions: vec![coinbase_with_value(10), spender].into_boxed_slice(),
        };
        let undo = undo_log_from_connected_block(&block, 5, |outpoint| {
            if *outpoint == spent {
                Ok(previous.clone())
            } else {
                Err(crate::error::ConsensusError::BlockValidation(
                    "missing prevout".into(),
                ))
            }
        })
        .unwrap();

        let coinbase_id = calculate_tx_id(&block.transactions[0]);
        let spend_id = calculate_tx_id(&block.transactions[1]);
        let mut set = UtxoSet::default();
        set.insert(
            OutPoint {
                hash: coinbase_id,
                index: 0,
            },
            Arc::new(UTXO {
                value: 10,
                script_pubkey: vec![0x51].into(),
                height: 5,
                is_coinbase: true,
            }),
        );
        set.insert(
            OutPoint {
                hash: spend_id,
                index: 0,
            },
            Arc::new(UTXO {
                value: 40,
                script_pubkey: vec![0x52].into(),
                height: 5,
                is_coinbase: false,
            }),
        );

        let restored = disconnect_block(&block, &undo, set, 5).unwrap();
        assert!(restored.contains_key(&spent));
        assert_eq!(restored.get(&spent).unwrap().value, 50);
        assert_eq!(restored.get(&spent).unwrap().height, 4);
        assert!(is_coinbase(&block.transactions[0]));
        assert!(!restored.contains_key(&OutPoint {
            hash: coinbase_id,
            index: 0,
        }));
        assert!(!restored.contains_key(&OutPoint {
            hash: spend_id,
            index: 0,
        }));
    }

    #[test]
    fn rebuild_refuses_a_spend_whose_output_is_unknown() {
        let spender = Transaction {
            version: 1,
            inputs: vec![TransactionInput {
                prevout: OutPoint {
                    hash: [9; 32],
                    index: 1,
                },
                script_sig: vec![].into(),
                sequence: 0xffffffff,
            }]
            .into(),
            outputs: vec![].into(),
            lock_time: 0,
        };
        let block = Block {
            header: BlockHeader {
                version: 1,
                prev_block_hash: [1; 32],
                merkle_root: [2; 32],
                timestamp: 1,
                bits: 0x207fffff,
                nonce: 0,
            },
            transactions: vec![spender].into_boxed_slice(),
        };
        let err = undo_log_from_connected_block(&block, 5, |_| {
            Err(crate::error::ConsensusError::BlockValidation(
                "missing prevout".into(),
            ))
        })
        .unwrap_err();
        let text = err.to_string();
        assert!(text.contains("missing prevout"), "{text}");
    }
}