watcher/src/core/blockchain.rs

use sha2::Digest;
use sha2::Sha256;

use vlogger::*;

use crate::core;
use crate::error::{ BlockchainError, TxError };
use std::collections::HashMap;
use std::io::Write;
use std::time::UNIX_EPOCH;

pub type Account = String;

#[derive(Debug)]
pub enum ValidationError {
    InvalidBlockHash,
    InvalidPreviousBlockHash
}

#[derive(Debug, Default)]
pub struct Blockchain {
    balances: std::collections::HashMap<Account, u32>,
    blocks: Vec<core::Block>,
    tx_mempool: Vec<core::Tx>,
}

#[allow(dead_code)]
impl Blockchain {
    pub fn open_account(&mut self, tx: core::Tx) -> Result<(), BlockchainError> {
        if !tx.is_new_account() {
            Err(BlockchainError::InvalidAccountCreation)
        } else {
            self.add(tx)
        }
    }

    pub fn add(&mut self, tx: core::Tx) -> Result<(), BlockchainError> {
        self.apply(&tx)?;
        self.tx_mempool.push(tx);
        Ok(())
    }

    pub fn hash_transaction(tx: &core::Tx) -> String {
        let mut hasher = Sha256::new();
        hasher.update(tx.to());
        hasher.update(tx.from());
        hasher.update(tx.value().to_be_bytes());
        hasher.update(tx.data());

        let res = hasher.finalize();
        hex::encode(res)
    }

    pub fn calculate_next_level(level: &[String]) -> Vec<String> {
        let mut next_level = Vec::new();

        for chunk in level.chunks(2) {
            let combined_hash = if chunk.len() == 2 {
                Self::hash_pair(&chunk[0], &chunk[1])
            } else {
                Self::hash_pair(&chunk[0], &chunk[0])
            };
            next_level.push(combined_hash);
        }
        next_level
    }

    pub fn calculate_merkle_root(tx: &[core::Tx]) -> String {
        let tx_hashes: Vec<String> = tx
            .iter()
            .map(|tx| Blockchain::hash_transaction(tx))
            .collect();

        if tx_hashes.is_empty() {
            return Blockchain::hash_data("");
        }

        if tx_hashes.len() == 1 {
            return tx_hashes[0].clone();
        }

        let mut current_level = tx_hashes.to_vec();

        while current_level.len() > 1 {
            current_level = Self::calculate_next_level(&current_level);
        }

        return current_level[0].clone();
    }

    fn hash_pair(left: &str, right: &str) -> String {
        let combined = format!("{}{}", left, right);
        Self::hash_data(&combined)
    }

    fn hash_data(data: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(data.as_bytes());
        hex::encode(hasher.finalize())
    }

    pub fn dump_blocks(&self, db_file: &mut std::fs::File) {
        let block_json = serde_json::to_string_pretty(&self.blocks).unwrap();
        db_file.write_all(&block_json.as_bytes()).unwrap();
    }

    pub fn create_block(&mut self) -> core::Block {
        let previous_hash = if self.blocks().len() > 0 {
            self.blocks().last().unwrap().head().block_hash()
        } else {
            ""
        };
        let merkle_root = Self::calculate_merkle_root(&self.tx_mempool);
        let timestamp = std::time::SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
        let nonce = 0;
        let mut new_head = core::BlockHeader {
            previous_hash: previous_hash.to_string(),
            merkle_root,
            timestamp,
            nonce,
            block_hash: "".to_string()
        };

        let mut block_hash = String::new();
        while !block_hash.starts_with("0") {
            new_head.nonce += 1;
            block_hash = calculate_block_hash(&new_head)
        }

        new_head.block_hash = block_hash;

        let new_block = core::Block::new(new_head, self.tx_mempool.clone());
        log!(DEBUG, "Created new Block {:#?}", new_block);
        self.blocks.push(new_block);
        self.blocks.last().unwrap().clone()
    }

    pub fn apply(&mut self, tx: &core::Tx) -> Result<(), BlockchainError> {
        self.tx_mempool.push(tx.clone());
        return Ok(());
        match tx.validate() {
            Ok(_) => {},
            Err(e) => return Err(BlockchainError::Tx(e))
        }

        if let Some(from_balance) = self.balances.get_mut(tx.from()) {
            if *from_balance > tx.value() {
                *from_balance -= tx.value();
            } else {
                return Err(BlockchainError::Tx(TxError::FromInsuffitientFonds))
            }
        } else {
            return Err(BlockchainError::Tx(TxError::UnknownAccount(tx.from().to_string())))
        }

        if let Some(to_balance) = self.balances.get_mut(&tx.to().to_string()) {
            *to_balance += tx.value()
        } else {
            if tx.is_new_account() {
                self.balances.insert(tx.to().to_string(), tx.value());
            } else {
                return Err(BlockchainError::Tx(TxError::UnknownAccount(tx.to().to_string())))
            }
        }

        Ok(())
    }

    pub fn new(balances: HashMap<Account, u32>, blocks: Vec<core::Block>, tx_mempool: Vec<core::Tx>) -> Blockchain {
        return Self {
            balances,
            blocks,
            tx_mempool
        }
    }

}

pub fn calculate_block_hash(head: &core::BlockHeader) -> String {
    let mut hasher = sha2::Sha256::new();

    hasher.update(head.nonce().to_be_bytes());
    hasher.update(head.previous_hash());
    hasher.update(head.timestamp().to_be_bytes());
    hasher.update(head.merkle_root());

    let res = hasher.finalize();
    hex::encode(res)
}

impl Blockchain {
    pub fn print_blocks(&self) {
        println!("Blocks List\n--------------");
        for (i, b) in self.blocks.iter().enumerate() {
            println!("Block #{i}\n{:#?}", b.head().block_hash);
        }
    }

    pub fn get_balances(&self) -> &std::collections::HashMap<String, u32> {
        &self.balances
    }

    pub fn blocks(&self) -> &[core::Block] {
        &self.blocks
    }

    pub fn add_block(&mut self, block: core::Block) {
        match self.validate_block(&block) {
            Ok(()) => self.blocks.push(block),
            Err(e) => match e {
                ValidationError::InvalidBlockHash => log!(ERROR, "Invalid Block Hash"),
                ValidationError::InvalidPreviousBlockHash => log!(ERROR, "Invalid Previos Block Hash")
            }
        }
    }

    fn validate_block(&self, block: &core::Block) -> Result<(), ValidationError>{
        let head = block.head();
        let hash = calculate_block_hash(block.head());
        if hash != head.block_hash() {
            return Err(ValidationError::InvalidBlockHash)
        }
        if let Some(prev_block) = self.blocks().last() {
            if head.previous_hash() != prev_block.head().block_hash() {
                return Err(ValidationError::InvalidPreviousBlockHash)
            }
        }
        Ok(())
    }

    fn validate_chain(&self) -> Result<(), ValidationError>{
        log!(INFO, "Validating Chain");
        let blocks = self.blocks();
        for block in blocks {
            let head = block.head();
            let hash = calculate_block_hash(block.head());

            if hash != head.block_hash() {
                return Err(ValidationError::InvalidBlockHash)
            }
        }
        Ok(())
    }

    pub fn build(blocks: Vec<core::Block>) -> Result<Blockchain, ValidationError> {
        log!(INFO, "Starting Chain Build from Genesis");
        let chain = Blockchain {
            blocks,
            balances: HashMap::new(),
            tx_mempool: vec![]
        };

        chain.validate_chain()?;
        Ok(chain)
    }
}