#![cfg_attr(not(feature = "std"), no_std)]

use frame_support::{
    ensure, traits::{Currency, Get},
    pallet_prelude::*,
};
use frame_system::pallet_prelude::*;

use sp_core::U256;
use sp_std::vec::Vec;
use sp_io::crypto::secp256k1_ecdsa_recover;
use sp_runtime::traits::{AtLeast32BitUnsigned, CheckedAdd, UniqueSaturatedInto};

use ghost_helpers::{
    SubstrateKeccakHasher, merkle_tree::verify_tree_proof,
    networks::{NetworkData, NetworkCurve, NetworkType, NetworkClaim},
};

use ghost_traits::{
    hashing::GhostHasher,
    networks::{
        NetworkDataBasicHandler, NetworkDataInspectHandler,
    },
};

pub use pallet::*;

mod weights;
pub use weights::WeightInfo;

mod benchmarking;
mod mock;
mod tests;

const MAX_POSSIBLE_PROOF_DEPTH: u32 = 256;

type TokenId = u128;
type EvmHash = sp_core::H256;
type EvmAddress = sp_core::H160;
type EvmSignature = [u8; 65];

#[derive(Default, Copy, Clone, Encode, Decode, RuntimeDebug, TypeInfo)]
pub struct ShareState<Balance>
where
	Balance: Default + Copy + Clone,
{
    total_shares: Balance,
    claimed_shares: Balance,
}

impl<Balance> ShareState<Balance>
where
	Balance: AtLeast32BitUnsigned + Default + Copy + Clone,
{
    pub fn new(total_shares: Balance) -> Self {
        let mut new_state = ShareState::default();
        new_state.total_shares = total_shares;
        new_state
    }

    pub fn try_increase_claimed(&mut self, increase_by: Balance) -> Option<Balance> {
        let new_claimed_shares = self.claimed_shares
            .saturating_add(increase_by);

        if new_claimed_shares > self.total_shares {
            return None;
        }

        self.claimed_shares = new_claimed_shares;

        Some(new_claimed_shares)
    }
}

#[derive(Default, Copy, Clone, Encode, Decode, RuntimeDebug, TypeInfo)]
pub struct NetworkShare<Balance>
where
	Balance: Default + Copy + Clone,
{
    proofs_size: u32,
    merkle_root: EvmHash,
    share_state: ShareState<Balance>
}

impl<Balance> NetworkShare<Balance>
where
	Balance: AtLeast32BitUnsigned + Default + Copy + Clone,
{
    pub fn new(
        total_shares: Balance,
        merkle_root: EvmHash,
        proofs_size: u32
    ) -> Self {
        let share_state = ShareState::new(total_shares);
        Self { merkle_root, share_state, proofs_size }
    }

    pub fn total_shares(&self) -> Balance {
        self.share_state.total_shares
    }

    pub fn claimed_shares(&self) -> Balance {
        self.share_state.claimed_shares
    }

    pub fn try_increase_claimed(&mut self, increase_by: Balance) -> Option<Balance> {
        self.share_state.try_increase_claimed(increase_by)
    }
}

#[derive(Default, Copy, Clone, Encode, Decode, RuntimeDebug, TypeInfo)]
pub struct MemberShare<Balance, NetworkId>
where
	Balance: AtLeast32BitUnsigned + Default + Copy,
	NetworkId: Default + Copy,
{
    initiated_network: Option<NetworkId>,
    activated_shares: Balance,
    locked_balance: Balance,
}

impl<Balance, NetworkId> MemberShare<Balance, NetworkId>
where
	Balance: AtLeast32BitUnsigned + Default + Copy,
	NetworkId: Default + Copy,
{
    pub fn increase_activated_shares(&mut self, extra_shares: Balance) {
        self.activated_shares = self.activated_shares
            .saturating_add(extra_shares);
    }

    pub fn initiate_on_network(&mut self, network_id: NetworkId) {
        self.initiated_network = Some(network_id);
    }

    pub fn lock_balance(&mut self, extra_balance: Balance) {
        self.locked_balance.saturating_add(extra_balance);
    }
}

#[derive(Encode, Decode, Clone, RuntimeDebug, TypeInfo, PartialEq, Eq)]
pub struct ClaimPackage<Balance>
where
	Balance: Default + Copy + Clone,
{
    pub index: u32,
    pub token_id: TokenId,
    pub shares: Balance,
    pub merkle_proof: BoundedVec<EvmHash, ConstU32<MAX_POSSIBLE_PROOF_DEPTH>>,
}

impl<Balance> ClaimPackage<Balance>
where
	Balance: Default + Copy + Clone + UniqueSaturatedInto<u128>,
{
    pub fn get_preimage<NetworkId>(
        &self,
        evm_address: &EvmAddress,
        network_id: NetworkId
    ) -> [u8; 128]
	where
		NetworkId: Default + Copy + Clone + UniqueSaturatedInto<u128>,
	{
        let mut token_id_bytes = [0u8; 32];
        U256::from(self.token_id as u128).to_big_endian(&mut token_id_bytes);

        let mut shares_bytes = [0u8; 32];
        let shares_u128: u128 = self.shares.unique_saturated_into();
        U256::from(shares_u128).to_big_endian(&mut shares_bytes);

        let mut address_bytes = [0u8; 32];
        address_bytes[12..32].copy_from_slice(evm_address.as_bytes());

        let mut network_id_bytes = [0u8; 32];
        let network_id_u128: u128 = network_id.unique_saturated_into();
        U256::from(network_id_u128).to_big_endian(&mut network_id_bytes);

        let mut preimage = [0u8; 128];
        preimage[0..32].copy_from_slice(&token_id_bytes);
        preimage[32..64].copy_from_slice(&shares_bytes);
        preimage[64..96].copy_from_slice(&address_bytes);
        preimage[96..128].copy_from_slice(&network_id_bytes);

        preimage
    }
}

pub type BalanceOf<T> =
    <<T as Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance;

pub type NetworkIdOf<T> =
	 <<T as Config>::NetworkDataHandler as NetworkDataBasicHandler>::NetworkId;

#[frame_support::pallet]
pub mod pallet {
    use super::*;

    #[pallet::pallet]
    #[pallet::without_storage_info]
    pub struct Pallet<T>(_);

    #[pallet::config]
    pub trait Config: frame_system::Config + core::fmt::Debug {
        type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;

        type Currency: Currency<Self::AccountId>;

        type NetworkDataHandler: NetworkDataInspectHandler<NetworkData>
			+ NetworkDataBasicHandler<NetworkCurve = NetworkCurve, NetworkType = NetworkType>;

        #[pallet::constant]
        type MinimumDonation: Get<u128>;

        #[pallet::constant]
        type MaxProofDepth: Get<u32>;

        type WeightInfo: WeightInfo;
    }

    #[pallet::event]
    #[pallet::generate_deposit(pub(super) fn deposit_event)]
    pub enum Event<T: Config> {
        SharesClaimed {
            who: T::AccountId,
            shares: BalanceOf<T>,
            network_id: NetworkIdOf<T>,
        }
    }

    #[pallet::error]
    pub enum Error<T> {
        PreimageAlreadyClaimed,
        BoundedToOtherNetwork,
        AuthorityNotGovernor,
        NetworkHasNoGenesis,
        InvalidMerkleProof,
        SharesOverflowed,
        BadSignature,
    }

    #[pallet::storage]
    #[pallet::getter(fn global_shares)]
    pub type GlobalShares<T: Config> = StorageValue<
        _,
        ShareState<BalanceOf<T>>,
        ValueQuery,
    >;

    #[pallet::storage]
    #[pallet::getter(fn network_shares)]
    pub type NetworkShares<T: Config> = StorageMap<
        _,
        Twox64Concat, NetworkIdOf<T>,
        NetworkShare<BalanceOf<T>>,
        ValueQuery,
    >;

    #[pallet::storage]
    #[pallet::getter(fn member_shares)]
    pub type MemberShares<T: Config> = CountedStorageMap<
        _,
        Blake2_256, T::AccountId,
        MemberShare<BalanceOf<T>, NetworkIdOf<T>>,
        ValueQuery,
    >;

    #[pallet::storage]
    #[pallet::getter(fn claimed_hashes)]
    pub type ClaimedHashes<T: Config> = StorageMap<
        _,
        Blake2_256, EvmHash,
        bool,
        ValueQuery,
    >;

    #[pallet::genesis_config]
    #[derive(frame_support::DefaultNoBound)]
    pub struct GenesisConfig<T: Config> {
        pub network_claims: Vec<NetworkClaim<NetworkIdOf<T>, BalanceOf<T>>>,
    }

    #[pallet::genesis_build]
    impl<T: Config> BuildGenesisConfig for GenesisConfig<T> {
        fn build(&self) {
            let mut global_shares: BalanceOf<T> = Default::default();

            self.network_claims
                .iter()
                .for_each(|c| {
                    assert!(
                        c.proof_size <= T::MaxProofDepth::get(),
                        "CRITICAL: constant T::MaxProofDepth is not enough for network {:?}",
                        c.network_id,
                    );

                    let network_state = NetworkShare::new(c.amount, c.merkle_root, c.proof_size);
                    NetworkShares::<T>::insert(c.network_id, network_state);

                    global_shares = global_shares.checked_add(&c.amount)
                        .expect("Global total shares overflowed");
                });

            let global_state = ShareState::new(global_shares);
            GlobalShares::<T>::put(global_state);
        }
    }

    #[pallet::call]
    impl<T: Config> Pallet<T> {
        #[pallet::call_index(0)]
        #[pallet::weight(<T as Config>::WeightInfo::claim())]
        pub fn claim(
            origin: OriginFor<T>,
            network_id: NetworkIdOf<T>,
            claim_package: ClaimPackage<BalanceOf<T>>,
            evm_address: EvmAddress,
            evm_signature: EvmSignature,
        ) -> DispatchResult {
            let who = ensure_signed(origin)?;

            let mut member_share = MemberShares::<T>::get(&who);
            let minimum_donation: BalanceOf<T> = T::MinimumDonation::get()
                .unique_saturated_into();

            ensure!(
                member_share.locked_balance >= minimum_donation,
                Error::<T>::AuthorityNotGovernor,
            );

            let correct_network_bound = member_share.initiated_network
                .map(|inner_network_id| inner_network_id == network_id)
                .unwrap_or(true);

            ensure!(correct_network_bound, Error::<T>::BoundedToOtherNetwork);

            let preimage = claim_package.get_preimage(&evm_address, network_id);
            let preimage_hash = SubstrateKeccakHasher::hash(&preimage);

            ensure!(
                !ClaimedHashes::<T>::contains_key(&preimage_hash),
                Error::<T>::PreimageAlreadyClaimed
            );

            Self::verify_signature(&who, &network_id, &preimage_hash, &evm_address, &evm_signature)?;
            Self::verify_merkle_proof(&network_id, &preimage, &claim_package)?;

            NetworkShares::<T>::try_mutate(&network_id, |state| -> DispatchResult {
                let package_proof_size = claim_package.merkle_proof.len() as u32;
                let proof_size_match = package_proof_size == state.proofs_size;
                ensure!(proof_size_match, Error::<T>::SharesOverflowed);

                let increased = state.try_increase_claimed(claim_package.shares);
                ensure!(increased.is_some(), Error::<T>::SharesOverflowed);
                Ok(())
            })?;

            GlobalShares::<T>::try_mutate(|state| -> DispatchResult {
                let increased = state.try_increase_claimed(claim_package.shares);
                ensure!(increased.is_some(), Error::<T>::SharesOverflowed);
                Ok(())
            })?;

            member_share.increase_activated_shares(claim_package.shares);
            member_share.initiate_on_network(network_id);

            MemberShares::<T>::insert(&who, member_share);
            ClaimedHashes::<T>::insert(&preimage_hash, true);

            Self::deposit_event(Event::<T>::SharesClaimed {
                shares: claim_package.shares,
                network_id,
                who,
            });

            Ok(())
        }
    }
}

impl<T: Config> Pallet<T> {
    fn to_ascii_hex(data: &[u8]) -> Vec<u8> {
        let mut r = Vec::with_capacity(data.len() * 2);
        let mut push_nibble = |n| r.push(if n < 10 { b'0' + n } else { b'a' - 10 + n });
        for &b in data.iter() {
            push_nibble(b / 16);
            push_nibble(b % 16);
        }
        r
    }

    fn ethereum_signable_message(
        receiver: &T::AccountId,
        network_id: &NetworkIdOf<T>,
        preimage_hash: &EvmHash,
    ) -> Vec<u8> {
        let prefix = b"GMV Claim:";
        let network_id_u64: u64 = (*network_id).unique_saturated_into();

        let receiver_ascii_hex = Self::to_ascii_hex(&receiver.encode());
        let preimage_ascii_hex = Self::to_ascii_hex(preimage_hash.as_ref());
        let network_ascii_hex = Self::to_ascii_hex(&network_id_u64.to_be_bytes());

        let what_length = receiver_ascii_hex.len()
            .saturating_add(preimage_ascii_hex.len())
            .saturating_add(network_ascii_hex.len())
            .saturating_add(10 + 11 + 12 + 11); // all 4 prefixes

        let mut what = Vec::with_capacity(what_length);
        what.extend_from_slice(prefix);

        what.extend_from_slice(b"\naccount:0x");
        what.extend_from_slice(&receiver_ascii_hex);

        what.extend_from_slice(b"\npreimage:0x");
        what.extend_from_slice(&preimage_ascii_hex);

        what.extend_from_slice(b"\nnetwork:0x");
        what.extend_from_slice(&network_ascii_hex);

        let mut l = what.len();
        let mut rev = Vec::new();
        while l > 0 {
            rev.push(b'0' + (l % 10) as u8);
            l /= 10;
        }

        let mut v = Vec::with_capacity(26 + rev.len() + what.len());
        v.extend_from_slice(b"\x19Ethereum Signed Message:\n");
        v.extend(rev.into_iter().rev());
        v.extend(what);

        v
    }

    fn verify_signature(
        who: &T::AccountId,
        network_id: &NetworkIdOf<T>,
        preimage_hash: &EvmHash,
        evm_address: &EvmAddress,
        evm_signature: &EvmSignature,
    ) -> DispatchResult {
        let message = Self::ethereum_signable_message(who, network_id, preimage_hash);
        let message_bytes = SubstrateKeccakHasher::hash(&message).to_fixed_bytes();

        let recovered_pubkey = secp256k1_ecdsa_recover(evm_signature, &message_bytes)
            .map_err(|_| Error::<T>::BadSignature)?;

        let recovered_pubkey_hash = SubstrateKeccakHasher::hash(&recovered_pubkey);
        let recovered_address = EvmAddress::from_slice(&recovered_pubkey_hash[12..32]);

        ensure!(recovered_address == *evm_address, Error::<T>::BadSignature);
        Ok(())
    }

    fn verify_merkle_proof(
        network_id: &NetworkIdOf<T>,
        preimage: &[u8],
        claim_package: &ClaimPackage<BalanceOf<T>>,
    ) -> DispatchResult {
        let network_state = NetworkShares::<T>::get(network_id);

        let is_merkle_proof_valid =
            verify_tree_proof::<SubstrateKeccakHasher, _>(
                preimage,
                &claim_package.merkle_proof,
                network_state.merkle_root,
                claim_package.index,
            );

        ensure!(is_merkle_proof_valid, Error::<T>::InvalidMerkleProof);
        Ok(())
    }
}