use bincode::Encode;
use sha2::{Digest, Sha256, digest::DynDigest};
use std::ops::Deref;

use serde::{Deserialize, Serialize};

use crate::{ZERO_PREFIX_AMOUNT, hash_chain::HashChain, transaction::Transaction};

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};
use tokio::sync::mpsc;
use tokio::task;

pub trait CryptoDigest {
    fn digest(&self, state: &mut impl DynDigest);
}

#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Encode)]
pub struct Block {
    pub prev_hash: Vec<u8>,
    pub trans: Transaction,
    pub pow: Option<u64>,
}

impl Block {
    fn new(prev_hash: Option<&[u8]>, trans: Transaction) -> Self {
        Self {
            prev_hash: prev_hash
                .unwrap_or(&[])
                .to_vec(),
            trans,
            pow: None,
        }
    }

    fn pref_zeros(hashed: &[u8]) -> Result<usize, usize> {
        hashed
            .iter()
            .try_fold(0, |nonce, x| {
                if nonce >= ZERO_PREFIX_AMOUNT {
                    Ok(nonce)
                } else if *x == 0 {
                    Ok(nonce + 1)
                } else {
                    Err(nonce)
                }
            })
    }

    pub async fn calc_pow(&self) -> u64 {
        let threads = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(4); // TODO: config

        let found = Arc::new(AtomicBool::new(false));
        let (tx, mut rx) = mpsc::unbounded_channel::<u64>();
        let prev_hash = Arc::new(
            self.prev_hash
                .clone(),
        );
        let trans = self.trans;
        for t in 0..threads {
            let tx = tx.clone();
            let prev_hash = Arc::clone(&prev_hash);
            let step = threads as u64;
            let found = found.clone();

            task::spawn_blocking(move || {
                let mut i = t as u64;
                while !found.load(Ordering::Relaxed) {
                    let mut hasher = Sha256::new();
                    Digest::update(&mut hasher, bincode::encode_to_vec((&*prev_hash, &trans, i), bincode::config::standard()).unwrap());
                    let hashed = hasher.finalize();
                    if Self::pref_zeros(&hashed).is_ok() {
                        if !found.swap(true, Ordering::Relaxed) {
                            _ = tx.send(i);
                        }
                        break;
                    }
                    i = i.wrapping_add(step);
                }
            });
        }

        // for h in handles {
        //     h.abort(); // h.join();
        // }

        rx.recv() // blocking
            .await
            .expect("miner dropped without sending")
    }

    pub async fn calc_set_pow(&mut self) {
        self.pow = Some(
            self.calc_pow()
                .await,
        );
    }

    pub fn verify_pow(&self) -> bool {
        if let Some(nonce) = self.pow {
            let mut hasher = Sha256::new();
            Digest::update(&mut hasher, bincode::encode_to_vec((&self.prev_hash, &self.trans, nonce), bincode::config::standard()).unwrap());
            Self::pref_zeros(&hasher.finalize()).is_ok()
        } else {
            false
        }
    }
}

impl CryptoDigest for Block {
    fn digest(&self, state: &mut impl DynDigest) {
        state.update(&bincode::encode_to_vec((&self.prev_hash, &self.trans, &self.pow), bincode::config::standard()).unwrap());
    }
}

#[derive(Debug, Default, Deserialize, Serialize, Clone, PartialEq)]
pub struct BlockChain {
    pub blocks: HashChain<Block>,
}

impl BlockChain {
    pub fn new(blocks: Vec<Block>) -> Self {
        Self {
            blocks: HashChain::new(blocks),
        }
    }

    pub fn new_block(&self, trans: Transaction) -> Block {
        let prev = self
            .blocks
            .last();
        let prev_hash = if let Some(b) = prev {
            Some(HashChain::hash_o_item(b))
        } else {
            None
        };
        Block::new(prev_hash.as_deref(), trans)
    }

    pub fn store(&mut self, block: Block) {
        // TODO: check config
        if block.verify_pow() {
            self.blocks
                .push(block)
        } else {
            eprint!("Err: Failed to store unverified block")
        }
    }
}

impl Deref for BlockChain {
    type Target = HashChain<Block>;

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