stratum-apps 0.9.0

use std::fmt;
use stratum_core::{
    bitcoin::{
        absolute::LockTime,
        script::ScriptBuf,
        transaction::{OutPoint, Sequence, Transaction, TxIn, TxOut, Version},
        witness::Witness,
    },
    template_distribution_sv2::CoinbaseOutputConstraintsOwned,
};

/// Offset added to the calculated max size to account for the additional space required by
/// Pay-to-Taproot (p2tr) addresses in solo mining scenarios.
///
/// Rationale: p2tr (Taproot) addresses produce the largest coinbase outputs among standard
/// address types. When the actual mining address is different from what was used to calculate
/// the constraints, we need extra space to accommodate the largest possible address type.
/// The value 43 bytes accounts for the difference between typical outputs and p2tr outputs.
const OFFSET_ADDITIONAL_SIZE: u32 = 43;

/// Offset added to the calculated max sigops to account for the additional signature operations
/// required by Pay-to-Public-Key-Hash (p2pkh) addresses in solo mining scenarios.
///
/// Rationale: p2pkh addresses have higher sigop counts compared to modern address types.
/// When the actual mining address differs from what was used to calculate constraints,
/// we need extra sigops budget to handle the worst-case scenario.
/// The value 4 sigops accounts for p2pkh's higher signature operation count.
const OFFSET_MAX_SIGOPS: u16 = 4;

/// Coinbase outputs needing more room than a `CoinbaseOutputConstraints` message can carry.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CoinbaseOutputConstraintsError {
    /// The outputs' serialized size, with any offset, is above `u32::MAX` bytes.
    SizeOutOfRange,
    /// The outputs' sigop cost, with any offset, is above `u16::MAX`.
    SigopsOutOfRange,
}

impl fmt::Display for CoinbaseOutputConstraintsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::SizeOutOfRange => write!(
                f,
                "coinbase outputs need more bytes than CoinbaseOutputConstraints can reserve"
            ),
            Self::SigopsOutOfRange => write!(
                f,
                "coinbase outputs need more sigops than CoinbaseOutputConstraints can reserve"
            ),
        }
    }
}

impl std::error::Error for CoinbaseOutputConstraintsError {}

/// Creates a CoinbaseOutputConstraints message from a list of coinbase outputs.
///
/// This function calculates the exact maximum additional size and sigops required
/// by the given coinbase outputs. No safety margins are added - the values reflect
/// precisely what the provided outputs require.
///
/// The size counts the outputs' serialized bytes alone. The CompactSize that counts the
/// coinbase's outputs grows from one byte to three past 252 outputs, and the Template Provider
/// reserves its worst case itself (sv2-spec 07, `CoinbaseOutputConstraints`), so counting it here
/// would reserve it twice.
///
/// # Arguments
/// * `coinbase_outputs` - List of transaction outputs that will be included in the coinbase
///
/// # Returns
/// CoinbaseOutputConstraints with exact size and sigops values based on the provided outputs
///
/// # Errors
/// Outputs whose size or sigop cost doesn't fit the message's fields return an error, never a
/// smaller reservation.
pub fn coinbase_output_constraints_message(
    coinbase_outputs: Vec<TxOut>,
) -> Result<CoinbaseOutputConstraintsOwned, CoinbaseOutputConstraintsError> {
    // calculate the max coinbase output size for CoinbaseOutputConstraints
    let max_size: usize = coinbase_outputs.iter().map(TxOut::size).sum();
    tracing::debug!(
        max_size,
        outputs_count = coinbase_outputs.len(),
        "Calculated max coinbase output size"
    );

    // this is used to calculate the sigops of the coinbase outputs
    // for CoinbaseOutputConstraints
    let dummy_coinbase = Transaction {
        version: Version::TWO,
        lock_time: LockTime::ZERO,
        input: vec![TxIn {
            previous_output: OutPoint::null(),
            script_sig: ScriptBuf::new(),
            sequence: Sequence::MAX,
            witness: Witness::from(vec![vec![0; 32]]),
        }],
        output: coinbase_outputs,
    };

    let max_sigops = dummy_coinbase.total_sigop_cost(|_| None);
    tracing::debug!(max_sigops, "Calculated max sigops for coinbase");

    Ok(CoinbaseOutputConstraintsOwned {
        coinbase_output_max_additional_size: u32::try_from(max_size)
            .map_err(|_| CoinbaseOutputConstraintsError::SizeOutOfRange)?,
        coinbase_output_max_additional_sigops: u16::try_from(max_sigops)
            .map_err(|_| CoinbaseOutputConstraintsError::SigopsOutOfRange)?,
    })
}

/// Creates a CoinbaseOutputConstraints message with safety margins (offsets) applied.
///
/// This function first calculates the exact constraints using
/// [`coinbase_output_constraints_message`], then adds safety margins to account for address type
/// variation in solo mining scenarios.
///
/// The offsets ensure that the coinbase can accommodate the largest possible address types
/// (p2tr for size, p2pkh for sigops) even when the actual mining address differs from what
/// was used to generate the coinbase outputs. This prevents validation failures when the
/// pool assigns a different address type than what the miner expected.
///
/// # Arguments
/// * `coinbase_outputs` - List of transaction outputs that will be included in the coinbase
///
/// # Returns
/// CoinbaseOutputConstraints with added safety margins for address type variation
///
/// # Offsets Applied
/// - **Size**: +43 bytes to accommodate Pay-to-Taproot (p2tr) addresses
/// - **Sigops**: +4 to accommodate Pay-to-Public-Key-Hash (p2pkh) addresses
///
/// # Errors
/// Outputs whose size or sigop cost, with the offsets added, doesn't fit the message's fields
/// return an error, never a smaller reservation.
pub fn coinbase_output_constraints_message_with_offset(
    coinbase_outputs: Vec<TxOut>,
) -> Result<CoinbaseOutputConstraintsOwned, CoinbaseOutputConstraintsError> {
    with_offset(coinbase_output_constraints_message(coinbase_outputs)?)
}

// `constraints` with the safety offsets added.
fn with_offset(
    constraints: CoinbaseOutputConstraintsOwned,
) -> Result<CoinbaseOutputConstraintsOwned, CoinbaseOutputConstraintsError> {
    Ok(CoinbaseOutputConstraintsOwned {
        coinbase_output_max_additional_size: constraints
            .coinbase_output_max_additional_size
            .checked_add(OFFSET_ADDITIONAL_SIZE)
            .ok_or(CoinbaseOutputConstraintsError::SizeOutOfRange)?,
        coinbase_output_max_additional_sigops: constraints
            .coinbase_output_max_additional_sigops
            .checked_add(OFFSET_MAX_SIGOPS)
            .ok_or(CoinbaseOutputConstraintsError::SigopsOutOfRange)?,
    })
}
#[cfg(test)]
mod tests {
    use super::*;
    use std::str::FromStr;
    use stratum_core::bitcoin::{Address, Amount, Network, TxOut};

    /// Informational test to showcase how the offset values were calculated.
    ///
    /// This test iterates through all standard Bitcoin address types and prints their
    /// sizes and sigop counts. Run with `--nocapture` to see the output.
    ///
    /// The offsets were determined by finding the worst-case address types:
    /// - **Size**: p2tr (Taproot) produces the largest coinbase output (43 bytes larger than
    ///   others)
    /// - **Sigops**: p2pkh has higher sigop count than modern address types (4 sigops)
    ///
    /// Since the pool cannot know in advance which address type the miner will use,
    /// we add these safety margins to ensure the coinbase can accommodate any address type.
    #[test]
    fn test_offset_values_rationale() {
        let addresses = vec![
            ("p2pkh", "19drg6CgjcvqFZSW5FLWdmqTBBeFLS5iC7"),
            ("p2sh", "18tRWCdi2Fc9CM57fUfmFK3ZC6cpGQeBkV"),
            ("p2wpkh", "bc1qwq787dzgj2w8hh58t4clr594y0cjgjashr0fz5"),
            ("p2wsh", "bc1qn04san36d0j76j0xksz2tesmtww7uf24j2x6v3"),
            (
                "p2tr",
                "bc1p8fltq0npm605tzl22gqewhy9dt5l25m7x67832vyhh9aem24rgdsgqwtpu",
            ),
        ];

        let mut max_size = 0u32;
        let mut max_sigops = 0u16;
        let mut max_size_type = "";
        let mut max_sigops_type = "";

        for (name, addr_str) in addresses {
            let addr = Address::from_str(addr_str)
                .unwrap()
                .require_network(Network::Bitcoin)
                .unwrap();

            let script = addr.script_pubkey();

            let coinbase_outputs = vec![TxOut {
                value: Amount::from_sat(50_0000_0000),
                script_pubkey: script.clone(),
            }];

            let dummy_coinbase = Transaction {
                version: Version::TWO,
                lock_time: LockTime::ZERO,
                input: vec![TxIn {
                    previous_output: OutPoint::null(),
                    script_sig: ScriptBuf::new(),
                    sequence: Sequence::MAX,
                    witness: Witness::from(vec![vec![0; 32]]),
                }],
                output: coinbase_outputs.clone(),
            };

            let size_of_current_address: u32 =
                coinbase_outputs.iter().map(|o| o.size() as u32).sum();
            let sigops_of_current_address = dummy_coinbase.total_sigop_cost(|_| None) as u16;

            if size_of_current_address > max_size {
                max_size = size_of_current_address;
                max_size_type = name;
            }

            if sigops_of_current_address > max_sigops {
                max_sigops = sigops_of_current_address;
                max_sigops_type = name;
            }

            println!("--- {}", name);
            println!("address: {}", addr);
            println!("script: {}", script.to_hex_string());
            println!(
                "max_size={} max_sigops={}",
                size_of_current_address, sigops_of_current_address
            );
        }

        println!("\n=== worst case summary ===");
        println!("largest output size: {} ({})", max_size, max_size_type);
        println!("largest sigops: {} ({})", max_sigops, max_sigops_type);
    }

    /// Tests that `coinbase_output_constraints_message` returns exact values without offsets.
    #[test]
    fn test_coinbase_output_constraints_message_exact() {
        let coinbase_outputs = vec![TxOut {
            value: Amount::from_sat(50_0000_0000),
            script_pubkey: ScriptBuf::from_hex(
                "0020aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
            )
            .unwrap(),
        }];

        let result = coinbase_output_constraints_message(coinbase_outputs.clone()).unwrap();

        let expected_size: u32 = coinbase_outputs.iter().map(|o| o.size() as u32).sum();
        let dummy_coinbase = Transaction {
            version: Version::TWO,
            lock_time: LockTime::ZERO,
            input: vec![TxIn {
                previous_output: OutPoint::null(),
                script_sig: ScriptBuf::new(),
                sequence: Sequence::MAX,
                witness: Witness::from(vec![vec![0; 32]]),
            }],
            output: coinbase_outputs,
        };
        let expected_sigops = dummy_coinbase.total_sigop_cost(|_| None) as u16;

        assert_eq!(result.coinbase_output_max_additional_size, expected_size);
        assert_eq!(
            result.coinbase_output_max_additional_sigops,
            expected_sigops
        );
    }

    /// Tests that `coinbase_output_constraints_message_with_offset` adds the correct offsets.
    #[test]
    fn test_coinbase_output_constraints_message_with_offset() {
        let coinbase_outputs = vec![TxOut {
            value: Amount::from_sat(50_0000_0000),
            script_pubkey: ScriptBuf::from_hex(
                "0020aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
            )
            .unwrap(),
        }];

        let result_with_offset =
            coinbase_output_constraints_message_with_offset(coinbase_outputs.clone()).unwrap();
        let result_exact = coinbase_output_constraints_message(coinbase_outputs).unwrap();

        assert_eq!(
            result_with_offset.coinbase_output_max_additional_size,
            result_exact.coinbase_output_max_additional_size + OFFSET_ADDITIONAL_SIZE
        );
        assert_eq!(
            result_with_offset.coinbase_output_max_additional_sigops,
            result_exact.coinbase_output_max_additional_sigops + OFFSET_MAX_SIGOPS
        );
    }

    /// Tests that offsets are applied correctly regardless of input.
    /// Uses a p2wpkh address to verify the offset logic.
    #[test]
    fn test_offset_values_are_applied_correctly() {
        let addr = Address::from_str("bc1qwq787dzgj2w8hh58t4clr594y0cjgjashr0fz5")
            .unwrap()
            .require_network(Network::Bitcoin)
            .unwrap();

        let coinbase_outputs = vec![TxOut {
            value: Amount::from_sat(50_0000_0000),
            script_pubkey: addr.script_pubkey(),
        }];

        let exact = coinbase_output_constraints_message(coinbase_outputs.clone()).unwrap();
        let with_offset =
            coinbase_output_constraints_message_with_offset(coinbase_outputs).unwrap();

        assert_eq!(
            with_offset.coinbase_output_max_additional_size,
            exact.coinbase_output_max_additional_size + OFFSET_ADDITIONAL_SIZE
        );
        assert_eq!(
            with_offset.coinbase_output_max_additional_sigops,
            exact.coinbase_output_max_additional_sigops + OFFSET_MAX_SIGOPS
        );
    }

    // One output whose script is `opcode` repeated `count` times.
    fn output_with_script(opcode: u8, count: usize) -> Vec<TxOut> {
        vec![TxOut {
            value: Amount::from_sat(0),
            script_pubkey: ScriptBuf::from_bytes(vec![opcode; count]),
        }]
    }

    const OP_CHECKSIG: u8 = 0xac;
    const OP_CHECKMULTISIG: u8 = 0xae;

    /// A sigop cost the message's 16-bit field can't carry is an error, never a smaller
    /// reservation: 1,000 `OP_CHECKMULTISIG` cost 80,000 sigops, and each legacy `OP_CHECKSIG`
    /// costs 4, so 16,383 of them fit the field and 16,384 don't.
    #[test]
    fn sigop_costs_the_message_cannot_carry_are_an_error() {
        assert_eq!(
            coinbase_output_constraints_message(output_with_script(OP_CHECKMULTISIG, 1_000)),
            Err(CoinbaseOutputConstraintsError::SigopsOutOfRange)
        );
        assert_eq!(
            coinbase_output_constraints_message(output_with_script(OP_CHECKSIG, 16_383))
                .unwrap()
                .coinbase_output_max_additional_sigops,
            65_532
        );
        assert_eq!(
            coinbase_output_constraints_message(output_with_script(OP_CHECKSIG, 16_384)),
            Err(CoinbaseOutputConstraintsError::SigopsOutOfRange)
        );
    }

    /// The offsets are added with checks: outputs costing 65,532 sigops fit the field exactly,
    /// but not with the 4-sigop offset, and either field one past its largest value with the
    /// offset is an error.
    #[test]
    fn offsets_past_the_message_fields_are_an_error() {
        assert_eq!(
            coinbase_output_constraints_message_with_offset(output_with_script(
                OP_CHECKSIG,
                16_383
            )),
            Err(CoinbaseOutputConstraintsError::SigopsOutOfRange)
        );

        let constraints = |size, sigops| CoinbaseOutputConstraintsOwned {
            coinbase_output_max_additional_size: size,
            coinbase_output_max_additional_sigops: sigops,
        };
        assert_eq!(
            with_offset(constraints(u32::MAX - OFFSET_ADDITIONAL_SIZE, 0)),
            Ok(constraints(u32::MAX, OFFSET_MAX_SIGOPS))
        );
        assert_eq!(
            with_offset(constraints(u32::MAX - OFFSET_ADDITIONAL_SIZE + 1, 0)),
            Err(CoinbaseOutputConstraintsError::SizeOutOfRange)
        );
        assert_eq!(
            with_offset(constraints(0, u16::MAX - OFFSET_MAX_SIGOPS)),
            Ok(constraints(OFFSET_ADDITIONAL_SIZE, u16::MAX))
        );
        assert_eq!(
            with_offset(constraints(0, u16::MAX - OFFSET_MAX_SIGOPS + 1)),
            Err(CoinbaseOutputConstraintsError::SigopsOutOfRange)
        );
    }

    /// The size is the outputs' serialized bytes alone, on both sides of the 253 outputs where
    /// the CompactSize counting them grows from one byte to three: the Template Provider reserves
    /// that count's worst case itself.
    #[test]
    fn the_size_leaves_the_output_count_to_the_template_provider() {
        for count in [252, 253] {
            let outputs = vec![
                TxOut {
                    value: Amount::from_sat(0),
                    script_pubkey: ScriptBuf::new(),
                };
                count
            ];
            let size: usize = outputs.iter().map(TxOut::size).sum();

            let constraints = coinbase_output_constraints_message(outputs).unwrap();

            assert_eq!(
                constraints.coinbase_output_max_additional_size as usize,
                size
            );
            assert_eq!(size, count * 9);
        }
    }
}