/*
    Copyright © 2023, ParallelChain Lab 
    Licensed under the Apache License, Version 2.0: http://www.apache.org/licenses/LICENSE-2.0
*/

//! Cryptography-related protocol types, including public addresses, secret keys, signatures, and hashes.
//! In protocol, we uses: 
//! - Ed25519 Signature 
//! - SHA256 Hash
//! - Merkle Tree Proof by using [rs_merkle]

use std::{fmt::Debug};

use borsh::BorshSerialize;
use rs_merkle::{MerkleTree, algorithms::Sha256};
use crate::{Serializable, Deserializable, consensus::Vote};

/// An Ed25519 signature. These are generated by external accounts to authorize transactions,
/// and by validators to create proposals and cast votes during consensus.
pub type Signature = [u8; 64];

/// An Ed25519 secret key. These are used to produce Ed25519 signatures. 
pub type SecretKey = [u8; 32];

/// PublicAddress is either:
/// - an Ed25519 public key representing an external account, or
/// - a contract address.
pub type PublicAddress = [u8; 32];

impl Serializable for PublicAddress {}
impl Deserializable for PublicAddress {}

/// A SHA256 hash. Used as block and transaction hashes, as well as to form Merkle tries.
pub type Sha256Hash = [u8; 32];

/// A space-efficient probabilistic data structure
pub type BloomFilter = [u8; 256];

// Computes the Merkle root hash of a vector of serializable data.
pub fn merkle_root<B: Serializable>(data: &Vec<B>) -> Sha256Hash {
    // null hash really isn't all 0s. There is no hash value for a tree without root. But here 
    // we use the 32-byte hash values to fill in the field definition inside data structures, for example, block header.
    if data.is_empty() {
        return [0; 32]
    }

    let leaves: Vec<[u8; 32]> = data
        .iter()
        .map(|datum| serialize_and_sha256::<B>(datum))
        .collect();
    let merkle_tree = MerkleTree::<Sha256>::from_leaves(&leaves);
    merkle_tree.root().unwrap()
}

/// Compute a Merkle Proof of inclusion of the leaf identified by `leaf_hash` inside `data`.
/// # Return value
/// If successful, returns a three tuple consisting of: 
/// 1. Leaf hashes.
/// 2. Root hashes.
/// 3. Proof.
pub fn merkle_proof<B: Serializable>(
    data: &Vec<B>,
    leaf_index: usize
) -> Result<(Vec<Sha256Hash>, Sha256Hash, Vec<u8>), LeafOutOfRangeError>  {
    if data.is_empty() {
        return Err(LeafOutOfRangeError);
    }

    let leaves: Vec<[u8; 32]> = data
        .iter()
        .map(|datum| serialize_and_sha256::<B>(datum))
        .collect();
    let merkle_tree = MerkleTree::<Sha256>::from_leaves(&leaves);

    Ok((leaves, merkle_tree.root().unwrap(), merkle_tree.proof(&[leaf_index]).to_bytes()))
}

/// Compute evidence which is a cryptographic hash over a pair of votes
pub fn evidence(
    signer: PublicAddress, 
    vote_1: &Vote, 
    vote_2: &Vote
) -> Sha256Hash {
    use sha2::{Sha256, Digest};
    let (vote_1, vote_2) =  match vote_1.chain_id.cmp(&vote_2.chain_id) {
        std::cmp::Ordering::Less => Some((vote_1, vote_2)),
        std::cmp::Ordering::Greater => Some((vote_2, vote_1)),
        std::cmp::Ordering::Equal => None 
    }.or_else(|| {
        match vote_1.view.cmp(&vote_2.view) {
            std::cmp::Ordering::Less => Some((vote_1, vote_2)),
            std::cmp::Ordering::Greater => Some((vote_2, vote_1)),
            std::cmp::Ordering::Equal => None 
        }
    }).or_else(||{
        match vote_1.block.cmp(&vote_2.block) {
            std::cmp::Ordering::Less => Some((vote_1, vote_2)),
            std::cmp::Ordering::Greater => Some((vote_2, vote_1)),
            std::cmp::Ordering::Equal => None,
        }
    }).or_else(||{
        match vote_1.phase.try_to_vec().unwrap().cmp(&vote_2.phase.try_to_vec().unwrap()) {
            std::cmp::Ordering::Less => Some((vote_1, vote_2)),
            std::cmp::Ordering::Greater => Some((vote_2, vote_1)),
            std::cmp::Ordering::Equal => None,
        }
    }).or_else(||{
        match vote_1.signature.cmp(&vote_2.signature) {
            std::cmp::Ordering::Less => Some((vote_1, vote_2)),
            std::cmp::Ordering::Greater => Some((vote_2, vote_1)),
            std::cmp::Ordering::Equal => None,
        }
    }).unwrap_or( (vote_1, vote_2) );  // basically identical

    Sha256::new()
        .chain_update(signer)
        .chain_update(vote_1.chain_id.to_le_bytes())
        .chain_update(vote_1.view.to_le_bytes())
        .chain_update(vote_1.block)
        .chain_update(vote_1.phase.try_to_vec().unwrap())
        .chain_update(vote_1.signature)
        .chain_update(vote_2.chain_id.to_le_bytes())
        .chain_update(vote_2.view.to_le_bytes())
        .chain_update(vote_2.block)
        .chain_update(vote_2.phase.try_to_vec().unwrap())
        .chain_update(vote_2.signature)
        .finalize().into()
}

pub struct LeafOutOfRangeError;

pub(crate) fn serialize_and_sha256<D: Serializable>(datum: &D) -> Sha256Hash {
        let r = <D as Serializable>::serialize(datum);
        sha256(&r)
}

/// Compute a Sha256 hash.
pub fn sha256(data: &[u8]) -> Sha256Hash {
    use sha2::{Sha256, Digest};
    let mut ret = Sha256::new();
    ret.update(data);
    ret.finalize().into()
}


pub trait AsKeyPair {
    fn as_keypair(&self) -> Result<ed25519_dalek::Keypair, CryptographicallyIncorrectTransactionError>;
}

impl AsKeyPair for [u8] {
    fn as_keypair(&self) -> Result<ed25519_dalek::Keypair, CryptographicallyIncorrectTransactionError> {
        ed25519_dalek::Keypair::from_bytes(self).map_err(|_| {CryptographicallyIncorrectTransactionError::InvalidKeypair})
    }
}

#[derive(Debug)]
pub enum CryptographicallyIncorrectTransactionError {
    InvalidFromAddress,
    InvalidKeypair,
    InvalidSignature,
    WrongSignature,
    WrongHash,
}