//! Push decoder for bitcoin blocks and transactions.

#![forbid(unsafe_code)]
#![allow(bare_trait_objects)]
#![allow(ellipsis_inclusive_range_patterns)]
#![warn(rustdoc::broken_intra_doc_links)]
#![warn(missing_docs)]
#![cfg_attr(all(not(feature = "std"), not(test)), no_std)]

#[cfg(not(any(feature = "std", feature = "no-std")))]
compile_error!("at least one of the `std` or `no-std` features must be enabled");

#[cfg(not(feature = "std"))]
extern crate core2;

use bitcoin::absolute::LockTime;
use bitcoin::hash_types::TxMerkleNode;
#[cfg(not(feature = "std"))]
use core2::io;
use merkle::IncrementalHasher;
#[cfg(feature = "std")]
use std::io;

extern crate alloc;
extern crate core;

/// Incremental merkle tree hasher
pub mod merkle;

use alloc::collections::VecDeque;
use alloc::vec::Vec;
use core::mem;
use core::ops::{Deref, DerefMut};

use bitcoin::consensus::{encode, Decodable, Encodable};
use bitcoin::hashes::{sha256::HashEngine, Hash, HashEngine as _};
use bitcoin::{
    OutPoint, ScriptBuf, Sequence, TxIn, TxOut, Txid,
    VarInt,
};
use bitcoin::blockdata::block::Header as BlockHeader;

use log::*;

/// Block decoder error
#[derive(Debug)]
pub enum Error {
    /// the input did not contain enough data to parse a block
    IncompleteData,
    /// the input contained invalid data
    ParseError,
    /// the input contained more data than the block
    TrailingData,
}

impl From<encode::Error> for Error {
    fn from(e: encode::Error) -> Self {
        debug!("parse error: {}", e);
        Error::ParseError
    }
}

impl From<crate::io::Error> for Error {
    fn from(e: crate::io::Error) -> Self {
        debug!("IO error: {}", e);
        Error::ParseError
    }
}

/// Block decoder listener.
pub trait Listener {
    /// called after a block header is parsed
    fn on_block_start(&mut self, header: &BlockHeader);
    /// called after a transaction header is parsed
    fn on_transaction_start(&mut self, version: i32);
    /// called after a transaction input is parsed
    fn on_transaction_input(&mut self, txin: &TxIn);
    /// called after a transaction output is parsed
    fn on_transaction_output(&mut self, txout: &TxOut);
    /// called after a transaction locktime is parsed at the end of the transaction
    fn on_transaction_end(&mut self, locktime: LockTime, txid: Txid);
    /// called after the complete block has been parsed
    fn on_block_end(&mut self);
}

#[derive(Debug, PartialEq)]
enum ParserState {
    BeforeHeader,
    ReadingTransactionHeader,
    // remaining inputs
    ReadingInputs(usize),
    // remaining inputs, prev outpoint, remaining script size
    ReadingInputScript(usize, OutPoint, usize),
    // remaining inputs, elements in input, bytes in element
    ReadingWitnesses(usize, usize, usize),
    BeforeOutputs,
    ReadingOutputs(usize),
    // remaining inputs, value, remaining script size
    ReadingOutputScript(usize, u64, usize),
    ReadingLockTime,
    FinishedBlock,
}

// Wrapper around a VecDeque<u8> that implements io::Read.
// Needed because Read was only added in 1.63 and core2 also doesn't have it.
struct Buffer(VecDeque<u8>);

impl Buffer {
    fn with_capacity(capacity: usize) -> Self {
        Self(VecDeque::with_capacity(capacity))
    }
}

impl Deref for Buffer {
    type Target = VecDeque<u8>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Buffer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl io::Read for Buffer {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let (ref mut front, _) = self.as_slices();
        let n = io::Read::read(front, buf)?;
        self.drain(..n);
        Ok(n)
    }
}

/// Push decoder for bitcoin blocks.
///
/// IMPORTANT: you must call `finish()` after the last byte of the block has been supplied,
/// in order to complete the block validation and ensure that there is no trailing data.
///
/// The decoder uses a limited amount of memory by streaming parse results
/// to a listener.
///
/// Scripts beyond `max_script_size` are replaced with the empty script (saves
/// 10 KB of peak memory usage).
///
/// The computed transactionm IDs cover the original script.
pub struct BlockDecoder {
    buffer: Buffer,
    buffer_capacity: usize,
    // scripts beyond this size are replaced with the empty script
    max_script_size: usize,
    parser_state: ParserState,
    // current script - reset after each input and output
    script: Option<Vec<u8>>,
    // number of transactions remaining in the block
    remaining_txs: usize,
    // number of inputs in the current transaction if the transaction has witnesses
    segwit_inputs: Option<usize>,
    // hasher for computing the current transaction's ID
    hasher: HashEngine,
    // merkle tree hasher
    merkle: IncrementalHasher,
    // merkle root from the header
    merkle_root: Option<TxMerkleNode>,
}

impl BlockDecoder {
    /// Create a new decoder with a max script size of 100 bytes
    pub fn new() -> Self {
        Self::new_with_capacity(100, 100)
    }

    /// Create a new decoder with a specific buffer capacity.
    /// must be at least 100 bytes.
    pub fn new_with_capacity(buffer_capacity: usize, max_script_size: usize) -> Self {
        assert!(buffer_capacity >= 100);
        let hasher = Txid::engine();
        Self {
            buffer: Buffer::with_capacity(buffer_capacity),
            buffer_capacity,
            max_script_size,
            parser_state: ParserState::BeforeHeader,
            script: None,
            remaining_txs: 0,
            segwit_inputs: None,
            hasher,
            merkle: IncrementalHasher::new(),
            merkle_root: None,
        }
    }

    /// Supply data to the decoder
    pub fn decode_next<L: Listener>(
        &mut self,
        mut data: &[u8],
        listener: &mut L,
    ) -> Result<(), Error> {
        while !data.is_empty() {
            let bytes_to_copy = usize::min(data.len(), self.buffer_capacity - self.buffer.len());
            trace!("copying {} bytes", bytes_to_copy);
            self.buffer.extend(&data[..bytes_to_copy]);
            data = &data[bytes_to_copy..];

            if !self.parse_step(listener)? {
                break;
            }
        }

        trace!("data is empty");
        // Attempt to completely parse any remaining data in the buffer
        self.parse_step(listener)?;
        trace!(
            "no progress possible at state {:?} len {}",
            self.parser_state,
            self.buffer.len()
        );
        Ok(())
    }

    /// Signal that no more data will be supplied, validate that the block is complete.
    pub fn finish(self) -> Result<(), Error> {
        assert_eq!(
            self.merkle_root,
            Some(self.merkle.finish()),
            "merkle root mismatch"
        );
        if self.parser_state != ParserState::FinishedBlock {
            Err(Error::IncompleteData)
        } else {
            Ok(())
        }
    }

    // parse as much as we can of the buffer and return true if we made progress
    fn parse_step<L: Listener>(&mut self, listener: &mut L) -> Result<bool, Error> {

        let initial_buffer_len = self.buffer.len();

        loop {
            trace!("state is {:?} len {}", self.parser_state, self.buffer.len());
            trace!("buffer {}", hex::encode(self.buffer.make_contiguous()));

            match self.parser_state {
                ParserState::BeforeHeader =>
                // less than 2^32 transactions - max 5 byte varint
                {
                    if self.buffer.len() >= 80 + 5 {
                        let header = BlockHeader::consensus_decode(&mut self.buffer)?;
                        listener.on_block_start(&header);
                        self.merkle_root = Some(header.merkle_root);
                        let tx_count = VarInt::consensus_decode(&mut self.buffer)?;
                        self.remaining_txs = tx_count.0 as usize;
                        // merkle tree must have at least one leaf
                        if self.remaining_txs == 0 {
                            return Err(Error::IncompleteData);
                        }
                        self.parser_state = ParserState::ReadingTransactionHeader;
                    } else {
                        break;
                    }
                }
                ParserState::ReadingTransactionHeader => {
                    // less than 2^32 outputs, max 5 byte varint
                    if self.buffer.len() >= 4 + 5 {
                        let version = i32::consensus_decode(&mut self.buffer)?;
                        version.consensus_encode(&mut self.hasher)?;

                        let mut input_count = VarInt::consensus_decode(&mut self.buffer)?;
                        if input_count.0 == 0 {
                            // segwit variant
                            let expected_one = VarInt::consensus_decode(&mut self.buffer)?;
                            if expected_one.0 != 1 {
                                return Err(Error::ParseError);
                            }

                            // now we have the real count
                            input_count = VarInt::consensus_decode(&mut self.buffer)?;
                            self.segwit_inputs = Some(input_count.0 as usize);
                        } else {
                            self.segwit_inputs = None;
                        }

                        input_count.consensus_encode(&mut self.hasher)?;

                        listener.on_transaction_start(version);
                        if input_count.0 > 0 {
                            self.parser_state = ParserState::ReadingInputs(input_count.0 as usize);
                        } else {
                            // not really allowed by bitcoind, but be lenient
                            self.parser_state = ParserState::BeforeOutputs;
                        }
                    } else {
                        break;
                    }
                }
                ParserState::ReadingInputs(remaining_inputs) => {
                    // the script length must not be longer than 10000 bytes, max 3 byte varint
                    if self.buffer.len() >= 36 + 3 {
                        let outpoint = OutPoint::consensus_decode(&mut self.buffer)?;
                        outpoint.consensus_encode(&mut self.hasher)?;

                        let script_len = VarInt::consensus_decode(&mut self.buffer)?;
                        script_len.consensus_encode(&mut self.hasher)?;
                        self.script = if script_len.0 > self.max_script_size as u64 {
                            None
                        } else {
                            Some(Vec::<u8>::with_capacity(script_len.0 as usize))
                        };
                        self.parser_state = ParserState::ReadingInputScript(
                            remaining_inputs,
                            outpoint,
                            script_len.0 as usize,
                        );
                    } else {
                        break;
                    }
                }
                ParserState::ReadingInputScript(
                    mut remaining_inputs,
                    outpoint,
                    mut remaining_script_len,
                ) => {
                    if self.buffer.is_empty() {
                        break;
                    }
                    let to_read = usize::min(remaining_script_len, self.buffer.len());

                    if let Some(ref mut script) = self.script {
                        script.extend(self.buffer.range(..to_read));
                    }

                    for byte in self.buffer.drain(..to_read) {
                        self.hasher.input(&[byte]);
                    }

                    remaining_script_len -= to_read;
                    self.parser_state = ParserState::ReadingInputScript(
                        remaining_inputs,
                        outpoint,
                        remaining_script_len,
                    );

                    if remaining_script_len == 0 {
                        if self.buffer.len() >= 4 {
                            let sequence = Sequence(u32::consensus_decode(&mut self.buffer)?);
                            sequence.consensus_encode(&mut self.hasher)?;
                            let script = self.script.take().unwrap_or(Vec::new());
                            let txin = TxIn {
                                previous_output: outpoint,
                                script_sig: ScriptBuf::from(script),
                                sequence,
                                witness: Default::default(),
                            };
                            listener.on_transaction_input(&txin);
                            remaining_inputs -= 1;

                            if remaining_inputs == 0 {
                                self.parser_state = ParserState::BeforeOutputs;
                            } else {
                                self.parser_state = ParserState::ReadingInputs(remaining_inputs);
                            }
                        } else {
                            // not enough data in buffer, wait for more
                            break;
                        }
                    }
                }
                ParserState::ReadingWitnesses(
                    mut remaining_inputs,
                    mut remaining_witnesses,
                    mut remaining_witness_len,
                ) => {
                    if remaining_witness_len > 0 {
                        if self.buffer.is_empty() {
                            break;
                        }
                        let to_read = usize::min(remaining_witness_len, self.buffer.len());
                        self.buffer.drain(..to_read);
                        remaining_witness_len -= to_read;
                        self.parser_state = ParserState::ReadingWitnesses(
                            remaining_inputs,
                            remaining_witnesses,
                            remaining_witness_len,
                        );
                    } else if remaining_witnesses > 0 {
                        if self.buffer.len() < 5 {
                            break;
                        }
                        let witness_length = VarInt::consensus_decode(&mut self.buffer)?.0 as usize;
                        remaining_witnesses -= 1;
                        self.parser_state = ParserState::ReadingWitnesses(
                            remaining_inputs,
                            remaining_witnesses,
                            witness_length,
                        );
                    } else if remaining_inputs > 0 {
                        if self.buffer.len() < 5 {
                            break;
                        }
                        let witnesses = VarInt::consensus_decode(&mut self.buffer)?.0 as usize;
                        remaining_inputs -= 1;
                        self.parser_state =
                            ParserState::ReadingWitnesses(remaining_inputs, witnesses, 0);
                    } else {
                        self.parser_state = ParserState::ReadingLockTime;
                    }
                }
                ParserState::BeforeOutputs => {
                    // less than 2^32 outputs
                    if self.buffer.len() >= 5 {
                        let output_count = VarInt::consensus_decode(&mut self.buffer)?;
                        output_count.consensus_encode(&mut self.hasher)?;
                        self.parser_state = ParserState::ReadingOutputs(output_count.0 as usize);
                    } else {
                        break;
                    }
                }
                ParserState::ReadingOutputs(remaining_outputs) => {
                    // the script length must not be longer than 10000 bytes, max 3 byte varint
                    if self.buffer.len() >= 8 + 3 {
                        let value = u64::consensus_decode(&mut self.buffer)?;
                        value.consensus_encode(&mut self.hasher)?;

                        let script_len = VarInt::consensus_decode(&mut self.buffer)?;
                        script_len.consensus_encode(&mut self.hasher)?;
                        self.script = if script_len.0 > self.max_script_size as u64 {
                            None
                        } else {
                            Some(Vec::with_capacity(script_len.0 as usize))
                        };
                        self.parser_state = ParserState::ReadingOutputScript(
                            remaining_outputs,
                            value,
                            script_len.0 as usize,
                        );
                    } else {
                        break;
                    }
                }
                ParserState::ReadingOutputScript(
                    mut remaining_outputs,
                    value,
                    mut remaining_script_len,
                ) => {
                    if self.buffer.is_empty() {
                        break;
                    }
                    let to_read = usize::min(remaining_script_len, self.buffer.len());
                    if let Some(ref mut script) = self.script {
                        script.extend(self.buffer.range(..to_read));
                    }
                    for byte in self.buffer.drain(..to_read) {
                        self.hasher.input(&[byte]);
                    }
                    remaining_script_len -= to_read;
                    if remaining_script_len > 0 {
                        self.parser_state = ParserState::ReadingOutputScript(
                            remaining_outputs,
                            value,
                            remaining_script_len,
                        );
                    } else {
                        let script = self.script.take().unwrap_or(Vec::new());
                        let txout = TxOut {
                            value,
                            script_pubkey: ScriptBuf::from(script),
                        };
                        listener.on_transaction_output(&txout);
                        remaining_outputs -= 1;

                        if remaining_outputs == 0 {
                            if let Some(segwit_inputs) = self.segwit_inputs {
                                self.parser_state =
                                    ParserState::ReadingWitnesses(segwit_inputs, 0, 0);
                            } else {
                                self.parser_state = ParserState::ReadingLockTime;
                            }
                        } else {
                            self.parser_state = ParserState::ReadingOutputs(remaining_outputs);
                        }
                    }
                }
                ParserState::ReadingLockTime => {
                    if self.buffer.len() >= 4 {
                        let locktime = LockTime::consensus_decode(&mut self.buffer)?;
                        locktime.consensus_encode(&mut self.hasher)?;
                        // this also resets the hasher
                        let engine = mem::take(&mut self.hasher);
                        let txid = Txid::from_engine(engine);
                        let txid_hash = txid.as_raw_hash();
                        self.merkle.add(TxMerkleNode::from_raw_hash(txid_hash.clone()));
                        listener.on_transaction_end(locktime, txid);
                        self.remaining_txs -= 1;
                        if self.remaining_txs == 0 {
                            self.parser_state = ParserState::FinishedBlock;
                            listener.on_block_end();
                        } else {
                            self.parser_state = ParserState::ReadingTransactionHeader;
                        }
                    } else {
                        break;
                    }
                }
                ParserState::FinishedBlock => {
                    // ensure that we have consumed all bytes from the buffer
                    if self.buffer.is_empty() {
                        break;
                    } else {
                        return Err(Error::TrailingData);
                    }
                }
            }
        }

        Ok(self.buffer.len() != initial_buffer_len)
    }
}

/// Test utilities
#[cfg(feature = "std")]
#[allow(missing_docs)]
pub mod test_util {
    use super::*;

    pub struct MockListener {
        pub block_headers: Vec<BlockHeader>,
        pub transaction_versions: Vec<i32>,
        pub transaction_inputs: Vec<TxIn>,
        pub transaction_outputs: Vec<TxOut>,
        pub transaction_locktimes: Vec<LockTime>,
        pub transaction_ids: Vec<Txid>,
        pub output_index: usize,
    }

    impl MockListener {
        pub fn new() -> Self {
            Self {
                block_headers: Vec::new(),
                transaction_versions: Vec::new(),
                transaction_inputs: Vec::new(),
                transaction_outputs: Vec::new(),
                transaction_locktimes: Vec::new(),
                transaction_ids: Vec::new(),
                output_index: 0,
            }
        }
    }

    impl Listener for MockListener {
        fn on_block_start(&mut self, header: &BlockHeader) {
            self.block_headers.push(header.clone());
        }

        fn on_transaction_start(&mut self, version: i32) {
            trace!("on_transaction_start({})", version);
            self.output_index = 0;
            self.transaction_versions.push(version);
        }

        fn on_transaction_input(&mut self, txin: &TxIn) {
            trace!("on_transaction_input: {:?}", txin);
            // fail fast
            assert_eq!(txin.sequence, Sequence::default());
            self.transaction_inputs.push(txin.clone());
        }

        fn on_transaction_output(&mut self, txout: &TxOut) {
            trace!("on_transaction_output: {:?}", txout);
            // fail fast
            assert_eq!(txout.value, self.output_index as u64);

            self.output_index += 1;
            self.transaction_outputs.push(txout.clone());
        }

        fn on_transaction_end(&mut self, locktime: LockTime, txid: Txid) {
            self.transaction_locktimes.push(locktime);
            self.transaction_ids.push(txid);
        }

        fn on_block_end(&mut self) {}
    }
}

#[cfg(test)]
mod tests {
    use alloc::vec::Vec;
    use bitcoin::block::Version;
    use bitcoin::consensus::{serialize, Encodable};
    use bitcoin::hashes::Hash;
    use bitcoin::{Block, BlockHash, CompactTarget, Transaction, Witness};
    use test_log::test;

    use super::test_util::MockListener;
    use super::*;

    #[test]
    fn test_decode_block_with_no_inputs() {
        let mut listener = MockListener::new();
        let mut decoder = BlockDecoder::new();
        // create a block with one transaction
        let mut block = Block {
            header: BlockHeader {
                version: Version::ONE,
                prev_blockhash: BlockHash::all_zeros(),
                merkle_root: TxMerkleNode::all_zeros(),
                time: 0,
                bits: CompactTarget::from_consensus(0),
                nonce: 0,
            },
            txdata: vec![Transaction {
                version: 1,
                lock_time: LockTime::from_consensus(0),
                input: vec![],
                output: vec![TxOut {
                    value: 0,
                    script_pubkey: ScriptBuf::from(vec![0x33, 0x44]),
                }],
            }],
        };
        block.header.merkle_root = block.compute_merkle_root().unwrap();
        let mut block_bytes = Vec::new();
        block.consensus_encode(&mut block_bytes).unwrap();
        trace!("block: {}\n{:#?}", hex::encode(block_bytes.clone()), block);
        decoder.decode_next(&block_bytes, &mut listener).unwrap();
        decoder.finish().unwrap();
        assert_eq!(listener.transaction_ids.len(), 1);
        assert_eq!(listener.transaction_inputs.len(), 0);
    }

    #[test]
    fn test_decode_block() {
        let mut listener = MockListener::new();
        let mut decoder = BlockDecoder::new();
        // create a block with one transaction
        let mut block = Block {
            header: BlockHeader {
                version: Version::ONE,
                prev_blockhash: BlockHash::all_zeros(),
                merkle_root: TxMerkleNode::all_zeros(),
                time: 0,
                bits: CompactTarget::from_consensus(0),
                nonce: 0,
            },
            txdata: vec![Transaction {
                version: 1,
                lock_time: LockTime::from_consensus(0),
                input: vec![TxIn {
                    previous_output: OutPoint {
                        txid: Txid::from_slice(&[0x33u8; 32]).unwrap(),
                        vout: 0x44,
                    },
                    script_sig: ScriptBuf::from(vec![0x11, 0x22]),
                    sequence: Default::default(),
                    witness: Default::default(),
                }],
                output: vec![TxOut {
                    value: 0,
                    script_pubkey: ScriptBuf::from(vec![0x33, 0x44]),
                }],
            }],
        };
        block.header.merkle_root = block.compute_merkle_root().unwrap();

        let mut block_bytes = Vec::new();
        block.consensus_encode(&mut block_bytes).unwrap();
        trace!("block: {}\n{:#?}", hex::encode(block_bytes.clone()), block);
        decoder.decode_next(&block_bytes, &mut listener).unwrap();
        decoder.finish().unwrap();
        assert_eq!(listener.block_headers.len(), 1);
        assert_eq!(listener.block_headers[0], block.header);

        assert_eq!(listener.transaction_versions.len(), 1);
        assert_eq!(listener.transaction_inputs.len(), 1);
        assert_eq!(
            listener.transaction_inputs[0].script_sig.as_bytes(),
            &[0x11, 0x22]
        );
        assert_eq!(listener.transaction_outputs.len(), 1);
        assert_eq!(
            listener.transaction_outputs[0].script_pubkey.as_bytes(),
            &[0x33, 0x44]
        );
        assert_eq!(listener.transaction_locktimes.len(), 1);

        assert_eq!(listener.transaction_ids.len(), 1);
        assert_eq!(listener.transaction_ids[0], block.txdata[0].txid());

        // parse a segwit block
        let mut listener = MockListener::new();
        let mut decoder = BlockDecoder::new();
        block.txdata[0].input[0].witness =
            Witness::from_slice(&vec![vec![0x11, 0x22], vec![], vec![0x33, 0x44]]);
        let mut block_bytes = Vec::new();
        block.consensus_encode(&mut block_bytes).unwrap();
        trace!("block: {}\n{:#?}", hex::encode(block_bytes.clone()), block);
        let ser = serialize(&block.txdata[0]);
        Transaction::consensus_decode(&mut &ser[..]).unwrap();
        trace!("tx: {}\n{:#?}", hex::encode(ser), block.txdata[0]);

        decoder.decode_next(&block_bytes, &mut listener).unwrap();
        decoder.finish().unwrap();
        assert_eq!(listener.transaction_ids.len(), 1);
        assert_eq!(listener.transaction_ids[0], block.txdata[0].txid());
    }
}