use crate::io::{Error, Write};
use alloc::collections::BTreeSet as OrderedSet;
use alloc::vec::Vec;
use bitcoin::consensus::{Decodable, Encodable};
use bitcoin::util::merkleblock::PartialMerkleTree;
use bitcoin::{Block, BlockHeader, OutPoint, Transaction, Txid};

/// An SPV proof that transactions are included in a block.
#[derive(Clone)]
pub struct SpvProof {
    /// The transactions in the proof.
    /// This may include additional transactions, if outputs from the requested ones are spent
    /// in the same block.
    /// These are in block order.
    pub txs: Vec<Transaction>,
    /// The consensus serialized merkle proof for the transactions, or empty if no transactions
    /// are included.
    pub proof: Option<PartialMerkleTree>,
}

impl Encodable for SpvProof {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, Error> {
        let mut len = 0;
        len += self.txs.consensus_encode(writer)?;
        len += 1;
        if let Some(proof) = &self.proof {
            true.consensus_encode(writer)?;
            len += proof.consensus_encode(writer)?;
        } else {
            false.consensus_encode(writer)?;
        }
        Ok(len)
    }
}

impl Decodable for SpvProof {
    fn consensus_decode<D: crate::io::Read + ?Sized>(
        reader: &mut D,
    ) -> Result<Self, bitcoin::consensus::encode::Error> {
        let txs = Decodable::consensus_decode(reader)?;
        let has_proof = Decodable::consensus_decode(reader)?;
        let proof = if has_proof {
            Some(Decodable::consensus_decode(reader)?)
        } else {
            None
        };
        Ok(SpvProof { txs, proof })
    }
}

impl SpvProof {
    /// Whether the proof proves nothing
    pub fn is_empty(&self) -> bool {
        self.txs.is_empty()
    }
    /// Build an SPV proof for a series of txids and outpoints that are being watched.
    /// - `txids` - transaction IDs that may be included in the block
    /// - `outpoints` - outpoints that may be spent in the block
    /// Returns None if there are no such spending transactions in the block.
    /// Also returns the spent outpoints.
    /// Also, returns the total set of unspent watched outpoints which includes unspent outputs
    /// of the transactions in the proof.  This can then be used to construct a proof of non-spend.
    pub fn build(
        block: &Block,
        txid_watches: &[Txid],
        outpoint_watches: &[OutPoint],
    ) -> (Self, Vec<OutPoint>, Vec<OutPoint>) {
        let watched_txids = OrderedSet::from_iter(txid_watches.iter());
        let mut watched_outpoints = OrderedSet::from_iter(outpoint_watches.iter().cloned());
        let mut spent_outpoints = OrderedSet::new();
        let mut txids = Vec::new();
        let mut matches = Vec::new();
        let mut matched_txs = Vec::new();
        for tx in block.txdata.iter() {
            let txid = tx.txid();
            txids.push(txid);
            let mut found_spent_outpoint = false;
            for input in tx.input.iter() {
                if watched_outpoints.remove(&input.previous_output) {
                    spent_outpoints.insert(input.previous_output);
                    found_spent_outpoint = true;
                }
            }
            if watched_txids.contains(&txid) || found_spent_outpoint {
                matches.push(true);
                matched_txs.push(tx.clone());
                // We need to watch the outputs of this transaction in case a subsequent
                // transaction in the block spends them.
                let additional_watches: Vec<OutPoint> = (0..tx.output.len() as u32)
                    .into_iter()
                    .map(|vout| OutPoint { txid, vout })
                    .collect();
                watched_outpoints.extend(additional_watches.into_iter());
            } else {
                matches.push(false);
            }
        }

        let proof = if matched_txs.is_empty() {
            None
        } else {
            Some(PartialMerkleTree::from_txids(&txids, &matches))
        };

        (
            Self {
                txs: matched_txs,
                proof,
            },
            spent_outpoints.into_iter().collect(),
            watched_outpoints.into_iter().collect(),
        )
    }

    /// Verify that the proof is valid for the given block
    pub fn verify(&self, header: &BlockHeader) -> bool {
        // Check SPV proof
        if let Some(txs_proof) = self.proof.as_ref() {
            let mut matches = Vec::new();
            let mut indexes = Vec::new();

            let extract = txs_proof.extract_matches(&mut matches, &mut indexes);
            let root = match extract {
                Ok(root) => root,
                Err(_) => return false,
            };
            if root != header.merkle_root {
                return false;
            }
            if matches.len() != self.txs.len() {
                return false;
            }
            for (tx, txid) in self.txs.iter().zip(matches) {
                if tx.txid() != txid {
                    return false;
                }
            }
        } else {
            if !self.txs.is_empty() {
                return false;
            }
        }
        return true;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::str::FromStr;

    use bitcoin::{
        Block, BlockHash, BlockHeader, OutPoint, PackedLockTime, Sequence, TxIn, TxMerkleNode,
        TxOut, Witness,
    };

    use bitcoin::hashes::Hash;

    fn make_tx(previous_outputs: Vec<OutPoint>, outputs: Vec<TxOut>) -> Transaction {
        Transaction {
            version: 2,
            lock_time: PackedLockTime::ZERO,
            input: previous_outputs
                .iter()
                .map(|previous_output| TxIn {
                    previous_output: *previous_output,
                    script_sig: Default::default(),
                    sequence: Sequence::ZERO,
                    witness: Witness::default(),
                })
                .collect(),
            output: outputs,
        }
    }

    fn make_blockheader() -> BlockHeader {
        BlockHeader {
            version: 0,
            prev_blockhash: BlockHash::all_zeros(),
            merkle_root: TxMerkleNode::from_str(
                "0377d5ba2c6e0f7aeaebb3caa6cd05b8b9b8ba60d0554e53b7a327ffdaa7433a",
            )
            .unwrap(),
            time: 0,
            bits: 0,
            nonce: 0,
        }
    }

    fn make_block() -> Block {
        Block {
            header: make_blockheader(),
            txdata: vec![
                make_tx(
                    // coinbase txid: 9310175d644aab7b9337ce806a4c7e27cdd815611085eab1a20c746a11742114
                    vec![OutPoint::from_str(
						"0000000000000000000000000000000000000000000000000000000000000000:4294967295").unwrap()],
                    vec![TxOut {
                        value: 5000002055,
                        script_pubkey: Default::default(),
                    }],
                ),
                make_tx(
                    // watched by txid_watch txid: 71ce38d3be3f07ac707d1c348bfa976f6d8060c28dce533914bcdc6b7e38d091
                    vec![OutPoint::from_str(
                        "7b2c3d17a43ac15757cc2b768d602d1a9333269802b8ee9fab375ea25a0509c8:0",
                    )
                    .unwrap()],
                    vec![TxOut {
                        value: 9993618,
                        script_pubkey: Default::default(),
                    }],
                ),
                make_tx(
                    // watched by watch txid: ea2e897dbe842b0a3645d7297070579f42ce234cfb1da25f9195f4259496a1a6
                    vec![OutPoint::from_str(
                        "5df6e0e2761359d30a8275058e299fcc0381534545f55cf43e41983f5d4c9456:1",
                    )
                    .unwrap()],
                    vec![TxOut {
                        value: 123456789,
                        script_pubkey: Default::default(),
                    }],
                ),
                make_tx(
                    // ignored txid: c8f94365a721c1637cacf44e734358d595dca8372a5d90e098482ed8f8d52cac
                    vec![OutPoint::from_str(
                        "f5864806e3565c34d1b41e716f72609d00b55ea5eac5b924c9719a842ef42206:1",
                    )
                    .unwrap()],
                    vec![TxOut {
                        value: 1000000000,
                        script_pubkey: Default::default(),
                    }],
                ),
                make_tx(
                    // ignored txid: 96b6df8e0cbb9919414c96a56a9b52e7299c9e394b46533c7dea2d14843a457b
                    vec![OutPoint::from_str(
                        "80b7d8a82d5d5bf92905b06f2014dd699e03837ca172e3a59d51426ebbe3e7f5:0",
                    )
                    .unwrap()],
                    vec![TxOut {
                        value: 2000000000,
                        script_pubkey: Default::default(),
                    }],
                ),
                make_tx(
                    // additional from watch txid: bfb031d4bb062a26561b828d686d0a30b2083d26463924060b638985348870cf
                    vec![OutPoint::from_str(
                        "ea2e897dbe842b0a3645d7297070579f42ce234cfb1da25f9195f4259496a1a6:0",
                    )
                    .unwrap()],
                    vec![TxOut {
                        value: 3000000000,
                        script_pubkey: Default::default(),
                    }],
                ),
                make_tx(
                    // additional from txid_watch txid: 17299ca25bd0dd92ef3e75f0aac5b04ad7477565e09e3bf4cc4fa2816f5d00d4
                    vec![OutPoint::from_str(
                        "71ce38d3be3f07ac707d1c348bfa976f6d8060c28dce533914bcdc6b7e38d091:0",
                    )
                    .unwrap()],
                    vec![TxOut {
                        value: 4000000000,
                        script_pubkey: Default::default(),
                    }],
                ),
            ],
        }
    }

    #[test]
    fn build_proof_with_empty_block_test() {
        let block = Block {
            header: make_blockheader(),
            txdata: vec![],
        };
        let proof = SpvProof::build(&block, &[], &[]).0;
        assert!(proof.verify(&block.header));
        assert!(proof.is_empty());
    }

    #[test]
    fn build_proof_with_empty_watches_test() {
        let block = make_block();
        let proof = SpvProof::build(&block, &[], &[]).0;
        assert!(proof.verify(&block.header));
        assert!(proof.is_empty());
    }

    #[test]
    fn bad_proof_test() {
        let block = make_block();
        let txid_watches = vec![block.txdata[1].txid()];
        let (proof, _spent, _unspent) = SpvProof::build(&block, &txid_watches, &[]);
        let header = BlockHeader {
            merkle_root: TxMerkleNode::all_zeros(),
            ..block.header
        };
        assert!(!proof.verify(&header));
        let mut proof1 = proof.clone();
        assert!(proof1.txs.pop().is_some());
        assert!(!proof1.verify(&block.header));

        let mut proof2 = proof.clone();
        proof2.txs[0].version = 5;
        assert!(!proof2.verify(&block.header));

        let mut proof3 = proof.clone();
        proof3.txs.push(block.txdata[0].clone());
        assert!(!proof3.verify(&block.header));

        let mut proof4 = proof.clone();
        proof4.proof = None;
        assert!(!proof4.verify(&block.header));
    }

    #[test]
    fn build_proof_with_txid_watch_test() {
        let block = make_block();
        let txid_watches = vec![block.txdata[1].txid()];
        let (proof, spent, unspent) = SpvProof::build(&block, &txid_watches, &[]);
        assert!(proof.verify(&block.header));
        // we got an additional tx because it spends an output of the watched tx
        assert_eq!(
            proof.txs,
            vec![block.txdata[1].clone(), block.txdata[6].clone()]
        );
        assert_eq!(
            spent.into_iter().collect::<Vec<_>>(),
            vec![OutPoint::new(block.txdata[1].txid(), 0)]
        );
        // the additional tx has an unspent outpoint, which was returned in build().1
        assert_eq!(
            unspent.into_iter().collect::<Vec<_>>(),
            vec![OutPoint {
                txid: block.txdata[6].txid(),
                vout: 0
            }]
        );
        let mut matches = Vec::new();
        let mut indexes = Vec::new();
        let root = proof
            .proof
            .unwrap()
            .extract_matches(&mut matches, &mut indexes)
            .unwrap();
        assert_eq!(root, block.header.merkle_root);
        assert_eq!(
            matches,
            vec![block.txdata[1].txid(), block.txdata[6].txid()]
        );
        assert_eq!(indexes, vec![1, 6]);
    }

    #[test]
    fn build_proof_with_outpoint_watch_test() {
        let block = make_block();
        let outpoint_watches = vec![block.txdata[2].input[0].previous_output];
        let (proof, spent, unspent) = SpvProof::build(&block, &[], &outpoint_watches);
        assert!(proof.verify(&block.header));
        // we got an additional tx because it spends an output of the watched tx
        assert_eq!(
            proof.txs,
            vec![block.txdata[2].clone(), block.txdata[5].clone()]
        );
        assert_eq!(
            spent.into_iter().collect::<Vec<_>>(),
            vec![
                outpoint_watches[0],
                OutPoint::new(block.txdata[2].txid(), 0),
            ]
        );
        // the additional tx has an unspent outpoint, which was returned in build().1
        assert_eq!(
            unspent.into_iter().collect::<Vec<_>>(),
            vec![OutPoint {
                txid: block.txdata[5].txid(),
                vout: 0
            }]
        );

        let mut matches = Vec::new();
        let mut indexes = Vec::new();
        let root = proof
            .proof
            .unwrap()
            .extract_matches(&mut matches, &mut indexes)
            .unwrap();
        assert_eq!(root, block.header.merkle_root);
        assert_eq!(
            matches,
            vec![block.txdata[2].txid(), block.txdata[5].txid()]
        );
        assert_eq!(indexes, vec![2, 5]);
    }
}
