rostrum 14.0.1

use crate::{
    chaindef::{Block, BlockHash, BlockHeader},
    undo::{BlockUndoer, DummyBlockUndoer},
};
use tokio::sync::RwLock;

use std::collections::HashMap;

/// A new header found, to be added to the chain at specific height
pub(crate) struct NewHeader {
    header: BlockHeader,
    hash: BlockHash,
    height: u64,
}

impl NewHeader {
    pub(crate) fn from_header_and_height(header: BlockHeader, height: u64) -> Self {
        let hash = header.block_hash();
        Self {
            header,
            hash,
            height,
        }
    }
}

/// Current blockchain headers' list
pub struct Chain {
    headers: RwLock<Vec<(BlockHash, BlockHeader)>>,
    heights: RwLock<HashMap<BlockHash, u64>>,
}

impl Chain {
    // create an empty chain
    pub fn new(genesis: Block) -> Self {
        let genesis_hash = genesis.block_hash();
        Self {
            headers: RwLock::new(vec![(genesis_hash, genesis.header)]),
            heights: RwLock::new(std::iter::once((genesis_hash, 0)).collect()), // genesis header @ zero height
        }
    }

    // Constructor for unit tests.
    #[cfg(bch)]
    pub fn new_regtest() -> Self {
        Self::new(crate::chaindef::regtest_genesis_block())
    }

    pub(crate) async fn drop_last_headers(&mut self, block_undoer: impl BlockUndoer, n: u64) {
        if n == 0 {
            return;
        }
        let new_height = self.height().await.saturating_sub(n);
        info!("Dropping {n} headers, undoing to block {new_height}");
        self.update(block_undoer, vec![], Some(new_height)).await
    }

    /// Load the chain from a collecion of headers, up to the given tip
    pub(crate) async fn load(&mut self, headers: Vec<BlockHeader>, tip: BlockHash) {
        let genesis_hash = self.headers.read().await[0].0;

        let mut header_map: HashMap<BlockHash, BlockHeader> =
            headers.into_iter().map(|h| (h.block_hash(), h)).collect();
        let mut blockhash = tip;
        let mut new_headers = vec![];
        while blockhash != genesis_hash {
            let header = match header_map.remove(&blockhash) {
                Some(header) => header,
                None => panic!("missing header {} while loading from DB", blockhash),
            };
            blockhash = header.prev_blockhash;
            new_headers.push(header);
        }
        info!("loading {} headers, tip={}", new_headers.len(), tip);
        let new_headers = new_headers.into_iter().rev(); // order by height
        let undoer = DummyBlockUndoer::default();
        self.update(
            undoer,
            new_headers
                .zip(1..)
                .map(|(header, height)| NewHeader::from_header_and_height(header, height))
                .collect(),
            None,
        )
        .await
    }

    /// Get the block hash at specified height (if exists)
    pub async fn get_block_hash(&self, height: usize) -> Option<BlockHash> {
        self.headers
            .read()
            .await
            .get(height)
            .map(|(hash, _header)| *hash)
    }

    /// Get the block header at specified height (if exists)
    pub(crate) async fn get_block_header(&self, height: usize) -> Option<BlockHeader> {
        self.headers
            .read()
            .await
            .get(height)
            .map(|(_hash, header)| header.clone())
    }

    /// Get the block height given the specified hash (if exists)
    pub(crate) async fn get_block_height(&self, blockhash: &BlockHash) -> Option<u64> {
        self.heights.read().await.get(blockhash).copied()
    }

    pub(crate) async fn get_block_height_and_header(
        &self,
        blockhash: &BlockHash,
    ) -> Option<(u64, BlockHeader)> {
        let headers = self.headers.read().await;
        let heights = self.heights.read().await;

        heights.get(blockhash).and_then(|height| {
            headers
                .get(*height as usize)
                .map(|(_, header)| (*height, header.clone()))
        })
    }

    pub(crate) async fn get_mtp(&self, height: usize) -> Option<u64> {
        let mut times = Vec::with_capacity(11);

        {
            let headers = self.headers.read().await;
            for h in height.saturating_sub(10)..=height {
                let (_, header) = headers.get(h)?;
                times.push(header.time as u64);
            }
        }

        if times.is_empty() {
            debug_assert!(false, "should always get mtp if block exists");
            return None;
        }

        times.sort_unstable();
        let mtp = times[times.len() / 2];

        Some(mtp)
    }

    /// Update the chain with a list of new headers (possibly a reorg)
    /// Tip height parameter is needed when chain shrinks and there are
    /// no new header (happens when invalidateblock is called).
    pub(crate) async fn update(
        &self,
        block_undoer: impl BlockUndoer,
        new_headers: Vec<NewHeader>,
        tip_height: Option<u64>,
    ) {
        let gensis = self.genesis_hash().await;
        let mut headers = self.headers.write().await;
        let mut heights = self.heights.write().await;

        let rewind = match new_headers.first() {
            Some(first) => first.height,
            None => tip_height.expect("empty new_headers and no tip height") + 1,
        };

        let tip_hash = tip_height.and_then(|h| headers.get(h as usize).map(|(hash, _)| *hash));
        let tip_hash = tip_hash.unwrap_or(gensis);

        let remove_headers: Vec<_> = headers.drain(rewind as usize..).collect();
        // undo in reverse order
        for (hash, _header) in remove_headers.into_iter().rev() {
            let height = heights.remove(&hash).expect("heights map missing entry");
            info!("Chain: Undoing block {} (height {})", hash, height);
            if let Err(e) = block_undoer.undo_block(&hash, height, &tip_hash).await {
                warn!(
                    "Failed to remove index entries for block {} (height {}): {}",
                    hash, height, e
                )
            };
        }

        if let Some(first_height) = new_headers.first().map(|h| h.height) {
            for (h, height) in new_headers.into_iter().zip(first_height..) {
                assert_eq!(h.height, height);
                assert_eq!(h.hash, h.header.block_hash());
                let existed = heights.insert(h.hash, h.height);
                assert!(existed.is_none());
                headers.push((h.hash, h.header));
            }
            info!(
                "chain updated: tip={}, height={}",
                headers.last().unwrap().0,
                headers.len() - 1
            );
        }
    }

    /// Best block hash
    pub async fn tip_hash(&self) -> BlockHash {
        self.headers.read().await.last().expect("empty chain").0
    }

    /// Best block header
    pub(crate) async fn tip(&self) -> BlockHeader {
        self.headers
            .read()
            .await
            .last()
            .expect("empty chain")
            .1
            .clone()
    }

    /**
     * Hash of the genesis block header
     */
    pub(crate) async fn genesis_hash(&self) -> BlockHash {
        self.headers.read().await.first().expect("empty chain").0
    }

    /// Number of blocks (excluding genesis block)
    pub(crate) async fn height(&self) -> u64 {
        (self.headers.read().await.len() - 1) as u64
    }

    /// If we have block
    pub async fn contains(&self, blockhash: &BlockHash) -> bool {
        self.heights.read().await.contains_key(blockhash)
    }
}

#[cfg(bch)]
#[cfg(test)]
mod tests {
    use super::*;
    use bitcoincash::consensus::deserialize;
    use bitcoincash::hashes::hex::{FromHex, ToHex};

    #[tokio::test]
    async fn test_genesis() {
        let regtest = Chain::new_regtest();
        assert_eq!(regtest.height().await, 0);
        assert_eq!(
            regtest.tip_hash().await.to_hex(),
            "0f9188f13cb7b2c71f2a335e3a4fc328bf5beb436012afca590b1a11466e2206"
        );
    }

    #[tokio::test]
    async fn test_updates() {
        let hex_headers = vec![
"0000002006226e46111a0b59caaf126043eb5bbf28c34f3a5e332a1fc7b2b73cf188910f1d14d3c7ff12d6adf494ebbcfba69baa915a066358b68a2b8c37126f74de396b1d61cc60ffff7f2000000000",
"00000020d700ae5d3c705702e0a5d9ababd22ded079f8a63b880b1866321d6bfcb028c3fc816efcf0e84ccafa1dda26be337f58d41b438170c357cda33a68af5550590bc1e61cc60ffff7f2004000000",
"00000020d13731bc59bc0989e06a5e7cab9843a4e17ad65c7ca47cd77f50dfd24f1f55793f7f342526aca9adb6ce8f33d8a07662c97d29d83b9e18117fb3eceecb2ab99b1e61cc60ffff7f2001000000",
"00000020a603def3e1255cadfb6df072946327c58b344f9bfb133e8e3e280d1c2d55b31c731a68f70219472864a7cb010cd53dc7e0f67e57f7d08b97e5e092b0c3942ad51f61cc60ffff7f2001000000",
"0000002041dd202b3b2edcdd3c8582117376347d48ff79ff97c95e5ac814820462012e785142dc360975b982ca43eecd14b4ba6f019041819d4fc5936255d7a2c45a96651f61cc60ffff7f2000000000",
"0000002072e297a2d6b633c44f3c9b1a340d06f3ce4e6bcd79ebd4c4ff1c249a77e1e37c59c7be1ca0964452e1735c0d2740f0d98a11445a6140c36b55770b5c0bcf801f1f61cc60ffff7f2000000000",
"000000200c9eb5889a8e924d1c4e8e79a716514579e41114ef37d72295df8869d6718e4ac5840f28de43ff25c7b9200aaf7873b20587c92827eaa61943484ca828bdd2e11f61cc60ffff7f2000000000",
"000000205873f322b333933e656b07881bb399dae61a6c0fa74188b5fb0e3dd71c9e2442f9e2f433f54466900407cf6a9f676913dd54aad977f7b05afcd6dcd81e98ee752061cc60ffff7f2004000000",
"00000020fd1120713506267f1dba2e1856ca1d4490077d261cde8d3e182677880df0d856bf94cfa5e189c85462813751ab4059643759ed319a81e0617113758f8adf67bc2061cc60ffff7f2000000000",
"000000200030d7f9c11ef35b89a0eefb9a5e449909339b5e7854d99804ea8d6a49bf900a0304d2e55fe0b6415949cff9bca0f88c0717884a5e5797509f89f856af93624a2061cc60ffff7f2002000000",
        ];
        let headers: Vec<BlockHeader> = hex_headers
            .iter()
            .map(|hex_header| deserialize(&Vec::from_hex(hex_header).unwrap()).unwrap())
            .collect();

        for chunk_size in 1..hex_headers.len() {
            let regtest = Chain::new_regtest();
            let mut height = 0;
            let mut tip = regtest.tip_hash().await;
            for chunk in headers.chunks(chunk_size) {
                let mut update = vec![];
                for header in chunk {
                    height += 1;
                    tip = header.block_hash();
                    update.push(NewHeader::from_header_and_height(*header, height))
                }
                regtest
                    .update(DummyBlockUndoer::default(), update, None)
                    .await;
                assert_eq!(regtest.tip_hash().await, tip);
                assert_eq!(regtest.height().await, height);
            }
            assert_eq!(
                regtest.tip_hash().await,
                headers.last().unwrap().block_hash()
            );
            assert_eq!(regtest.height().await, headers.len() as u64);
        }

        // test loading from a list of headers and tip
        let mut regtest = Chain::new_regtest();
        regtest
            .load(headers.clone(), headers.last().unwrap().block_hash())
            .await;
        assert_eq!(regtest.height().await, headers.len() as u64);

        // test getters
        for (header, height) in headers.iter().zip(1usize..) {
            assert_eq!(regtest.get_block_header(height).await, Some(*header));
            assert_eq!(
                regtest.get_block_hash(height).await,
                Some(header.block_hash())
            );
            assert_eq!(
                regtest.get_block_height(&header.block_hash()).await,
                Some(height as u64)
            );
        }

        // test chain shortening
        for i in (0..=headers.len()).rev() {
            let hash = regtest.get_block_hash(i).await.unwrap();
            assert_eq!(regtest.get_block_height(&hash).await, Some(i as u64));
            assert_eq!(regtest.height().await, i as u64);
            assert_eq!(regtest.tip_hash().await, hash);
            regtest
                .drop_last_headers(DummyBlockUndoer::new(true), 1)
                .await;
        }
        assert_eq!(regtest.height().await, 0);
        assert_eq!(
            regtest.tip_hash().await.to_hex(),
            "0f9188f13cb7b2c71f2a335e3a4fc328bf5beb436012afca590b1a11466e2206"
        );

        regtest
            .drop_last_headers(DummyBlockUndoer::new(true), 1)
            .await;
        assert_eq!(regtest.height().await, 0);
        assert_eq!(
            regtest.tip_hash().await.to_hex(),
            "0f9188f13cb7b2c71f2a335e3a4fc328bf5beb436012afca590b1a11466e2206"
        );

        // test reorg
        let mut regtest = Chain::new_regtest();
        regtest
            .load(headers.clone(), headers.last().unwrap().block_hash())
            .await;
        let height = regtest.height().await;

        let new_header: BlockHeader = deserialize(&Vec::from_hex("000000200030d7f9c11ef35b89a0eefb9a5e449909339b5e7854d99804ea8d6a49bf900a0304d2e55fe0b6415949cff9bca0f88c0717884a5e5797509f89f856af93624a7a6bcc60ffff7f2000000000").unwrap()).unwrap();
        regtest
            .update(
                DummyBlockUndoer::new(true),
                vec![NewHeader::from_header_and_height(new_header, height)],
                None,
            )
            .await;
        assert_eq!(regtest.height().await, height);
        assert_eq!(
            regtest.tip_hash().await.to_hex(),
            "0e16637fe0700a7c52e9a6eaa58bd6ac7202652103be8f778680c66f51ad2e9b"
        );
    }
}