use super::{Apply, ApplyBackend, Backend, Basic, Log};
use crate::core::utils::{U256_ONE, U256_ZERO};
use crate::prelude::*;
use primitive_types::{H160, H256, U256};

/// Vicinity value of a memory backend.
#[derive(Clone, Debug, Eq, PartialEq)]
#[cfg_attr(
    feature = "with-codec",
    derive(scale_codec::Encode, scale_codec::Decode, scale_info::TypeInfo)
)]
#[cfg_attr(feature = "with-serde", derive(serde::Serialize, serde::Deserialize))]
pub struct MemoryVicinity {
    /// Gas price.
    pub gas_price: U256,
    /// Effective gas price.
    pub effective_gas_price: U256,
    /// Origin.
    pub origin: H160,
    /// Chain ID.
    pub chain_id: U256,
    /// Environmental block hashes.
    pub block_hashes: Vec<H256>,
    /// Environmental block number.
    pub block_number: U256,
    /// Environmental coinbase.
    pub block_coinbase: H160,
    /// Environmental block timestamp.
    pub block_timestamp: U256,
    /// Environmental block difficulty.
    pub block_difficulty: U256,
    /// Environmental block gas limit.
    pub block_gas_limit: U256,
    /// Environmental base fee per gas.
    ///
    pub block_base_fee_per_gas: U256,
    /// Environmental randomness.
    ///
    /// In Ethereum, this is the randomness beacon provided by the beacon
    /// chain and is only enabled post Merge.
    pub block_randomness: Option<H256>,
    /// EIP-4844
    pub blob_gas_price: Option<u128>,
    /// EIP-4844
    pub blob_hashes: Vec<U256>,
}

/// Account information of a memory backend.
#[derive(Default, Clone, Debug, Eq, PartialEq)]
#[cfg_attr(
    feature = "with-codec",
    derive(scale_codec::Encode, scale_codec::Decode, scale_info::TypeInfo)
)]
#[cfg_attr(feature = "with-serde", derive(serde::Serialize, serde::Deserialize))]
pub struct MemoryAccount {
    /// Account nonce.
    pub nonce: U256,
    /// Account balance.
    pub balance: U256,
    /// Full account storage.
    pub storage: BTreeMap<H256, H256>,
    /// Account code.
    pub code: Vec<u8>,
}

/// Memory backend, storing all state values in a `BTreeMap` in memory.
#[derive(Clone, Debug)]
pub struct MemoryBackend<'vicinity> {
    vicinity: &'vicinity MemoryVicinity,
    state: BTreeMap<H160, MemoryAccount>,
    logs: Vec<Log>,
}

impl<'vicinity> MemoryBackend<'vicinity> {
    /// Create a new memory backend.
    #[must_use]
    pub const fn new(
        vicinity: &'vicinity MemoryVicinity,
        state: BTreeMap<H160, MemoryAccount>,
    ) -> Self {
        Self {
            vicinity,
            state,
            logs: Vec::new(),
        }
    }

    /// Get the underlying `BTreeMap` storing the state.
    #[must_use]
    pub const fn state(&self) -> &BTreeMap<H160, MemoryAccount> {
        &self.state
    }

    /// Get a mutable reference to the underlying `BTreeMap` storing the state.
    pub fn state_mut(&mut self) -> &mut BTreeMap<H160, MemoryAccount> {
        &mut self.state
    }
}

impl Backend for MemoryBackend<'_> {
    #[allow(clippy::misnamed_getters)]
    fn gas_price(&self) -> U256 {
        self.vicinity.effective_gas_price
    }
    fn origin(&self) -> H160 {
        self.vicinity.origin
    }
    fn block_hash(&self, number: U256) -> H256 {
        if number >= self.vicinity.block_number
            || self.vicinity.block_number - number - U256_ONE
                >= U256::from(self.vicinity.block_hashes.len())
        {
            H256::default()
        } else {
            let index = (self.vicinity.block_number - number - U256_ONE).as_usize();
            self.vicinity.block_hashes[index]
        }
    }
    fn block_number(&self) -> U256 {
        self.vicinity.block_number
    }
    fn block_coinbase(&self) -> H160 {
        self.vicinity.block_coinbase
    }
    fn block_timestamp(&self) -> U256 {
        self.vicinity.block_timestamp
    }
    fn block_difficulty(&self) -> U256 {
        self.vicinity.block_difficulty
    }
    fn block_randomness(&self) -> Option<H256> {
        self.vicinity.block_randomness
    }
    fn block_gas_limit(&self) -> U256 {
        self.vicinity.block_gas_limit
    }
    fn block_base_fee_per_gas(&self) -> U256 {
        self.vicinity.block_base_fee_per_gas
    }

    fn chain_id(&self) -> U256 {
        self.vicinity.chain_id
    }

    fn exists(&self, address: H160) -> bool {
        self.state.contains_key(&address)
    }

    fn basic(&self, address: H160) -> Basic {
        self.state
            .get(&address)
            .map(|a| Basic {
                balance: a.balance,
                nonce: a.nonce,
            })
            .unwrap_or_default()
    }

    fn code(&self, address: H160) -> Vec<u8> {
        self.state
            .get(&address)
            .map(|v| v.code.clone())
            .unwrap_or_default()
    }

    fn storage(&self, address: H160, index: H256) -> H256 {
        self.state
            .get(&address)
            .map(|v| v.storage.get(&index).copied().unwrap_or_default())
            .unwrap_or_default()
    }

    fn is_empty_storage(&self, address: H160) -> bool {
        self.state
            .get(&address)
            .map_or(true, |v| v.storage.is_empty())
    }

    fn original_storage(&self, address: H160, index: H256) -> Option<H256> {
        Some(self.storage(address, index))
    }
    fn blob_gas_price(&self) -> Option<u128> {
        self.vicinity.blob_gas_price
    }
    fn get_blob_hash(&self, index: usize) -> Option<U256> {
        self.vicinity.blob_hashes.get(index).copied()
    }
}

impl ApplyBackend for MemoryBackend<'_> {
    fn apply<A, I, L>(&mut self, values: A, logs: L, delete_empty: bool)
    where
        A: IntoIterator<Item = Apply<I>>,
        I: IntoIterator<Item = (H256, H256)>,
        L: IntoIterator<Item = Log>,
    {
        for apply in values {
            match apply {
                Apply::Modify {
                    address,
                    basic,
                    code,
                    storage,
                    reset_storage,
                } => {
                    let is_empty = {
                        let account = self.state.entry(address).or_default();
                        account.balance = basic.balance;
                        account.nonce = basic.nonce;
                        if let Some(code) = code {
                            account.code = code;
                        }

                        if reset_storage {
                            account.storage = BTreeMap::new();
                        }

                        let zeros = account
                            .storage
                            .iter()
                            .filter(|(_, v)| v == &&H256::default())
                            .map(|(k, _)| *k)
                            .collect::<Vec<H256>>();

                        for zero in zeros {
                            account.storage.remove(&zero);
                        }

                        for (index, value) in storage {
                            if value == H256::default() {
                                account.storage.remove(&index);
                            } else {
                                account.storage.insert(index, value);
                            }
                        }

                        account.balance == U256_ZERO
                            && account.nonce == U256_ZERO
                            && account.code.is_empty()
                    };

                    if is_empty && delete_empty {
                        self.state.remove(&address);
                    }
                }
                Apply::Delete { address } => {
                    self.state.remove(&address);
                }
            }
        }

        for log in logs {
            self.logs.push(log);
        }
    }
}
