use alloc::{
    string::{String, ToString},
    vec::Vec,
};
use core::fmt;
use super::{
    get_account_seed, AccountError, ByteReader, Deserializable, DeserializationError, Digest, Felt,
    Hasher, Serializable, Word, ZERO,
};
use crate::{crypto::merkle::LeafIndex, utils::hex_to_bytes, ACCOUNT_TREE_DEPTH};
pub const ACCOUNT_STORAGE_MASK_SHIFT: u64 = 62;
pub const ACCOUNT_STORAGE_MASK: u64 = 0b11 << ACCOUNT_STORAGE_MASK_SHIFT;
pub const ACCOUNT_TYPE_MASK_SHIFT: u64 = 60;
pub const ACCOUNT_TYPE_MASK: u64 = 0b11 << ACCOUNT_TYPE_MASK_SHIFT;
pub const ACCOUNT_ISFAUCET_MASK: u64 = 0b10 << ACCOUNT_TYPE_MASK_SHIFT;
pub const FUNGIBLE_FAUCET: u64 = 0b10;
pub const NON_FUNGIBLE_FAUCET: u64 = 0b11;
pub const REGULAR_ACCOUNT_IMMUTABLE_CODE: u64 = 0b00;
pub const REGULAR_ACCOUNT_UPDATABLE_CODE: u64 = 0b01;
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u64)]
pub enum AccountType {
    FungibleFaucet = FUNGIBLE_FAUCET,
    NonFungibleFaucet = NON_FUNGIBLE_FAUCET,
    RegularAccountImmutableCode = REGULAR_ACCOUNT_IMMUTABLE_CODE,
    RegularAccountUpdatableCode = REGULAR_ACCOUNT_UPDATABLE_CODE,
}
pub const fn account_type_from_u64(value: u64) -> AccountType {
    debug_assert!(
        ACCOUNT_TYPE_MASK.count_ones() == 2,
        "This method assumes there are only 2bits in the mask"
    );
    let bits = (value & ACCOUNT_TYPE_MASK) >> ACCOUNT_TYPE_MASK_SHIFT;
    match bits {
        REGULAR_ACCOUNT_UPDATABLE_CODE => AccountType::RegularAccountUpdatableCode,
        REGULAR_ACCOUNT_IMMUTABLE_CODE => AccountType::RegularAccountImmutableCode,
        FUNGIBLE_FAUCET => AccountType::FungibleFaucet,
        NON_FUNGIBLE_FAUCET => AccountType::NonFungibleFaucet,
        _ => {
            unreachable!()
        },
    }
}
impl From<u64> for AccountType {
    fn from(value: u64) -> Self {
        account_type_from_u64(value)
    }
}
pub const ON_CHAIN: u64 = 0b00;
pub const OFF_CHAIN: u64 = 0b10;
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u64)]
pub enum AccountStorageType {
    OnChain = ON_CHAIN,
    OffChain = OFF_CHAIN,
}
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
#[cfg_attr(feature = "serde", serde(transparent))]
pub struct AccountId(Felt);
impl AccountId {
    #[cfg(not(any(feature = "testing", test)))]
    pub const REGULAR_ACCOUNT_SEED_DIGEST_MIN_TRAILING_ZEROS: u32 = 23;
    #[cfg(not(any(feature = "testing", test)))]
    pub const FAUCET_SEED_DIGEST_MIN_TRAILING_ZEROS: u32 = 31;
    #[cfg(any(feature = "testing", test))]
    pub const REGULAR_ACCOUNT_SEED_DIGEST_MIN_TRAILING_ZEROS: u32 = 5;
    #[cfg(any(feature = "testing", test))]
    pub const FAUCET_SEED_DIGEST_MIN_TRAILING_ZEROS: u32 = 6;
    pub const MIN_ACCOUNT_ONES: u32 = 5;
    pub fn new(
        seed: Word,
        code_commitment: Digest,
        storage_root: Digest,
    ) -> Result<Self, AccountError> {
        let seed_digest = compute_digest(seed, code_commitment, storage_root);
        Self::validate_seed_digest(&seed_digest)?;
        seed_digest[0].try_into()
    }
    pub fn new_unchecked(value: Felt) -> Self {
        Self(value)
    }
    #[cfg(any(feature = "testing", test))]
    pub fn new_dummy(init_seed: [u8; 32], account_type: AccountType) -> Self {
        let code_commitment = Digest::default();
        let storage_root = Digest::default();
        let seed = get_account_seed(
            init_seed,
            account_type,
            AccountStorageType::OnChain,
            code_commitment,
            storage_root,
        )
        .unwrap();
        Self::new(seed, code_commitment, storage_root).unwrap()
    }
    pub const fn account_type(&self) -> AccountType {
        account_type_from_u64(self.0.as_int())
    }
    pub fn is_faucet(&self) -> bool {
        matches!(
            self.account_type(),
            AccountType::FungibleFaucet | AccountType::NonFungibleFaucet
        )
    }
    pub fn is_regular_account(&self) -> bool {
        is_regular_account(self.0.as_int())
    }
    pub fn storage_type(&self) -> AccountStorageType {
        let bits = (self.0.as_int() & ACCOUNT_STORAGE_MASK) >> ACCOUNT_STORAGE_MASK_SHIFT;
        match bits {
            ON_CHAIN => AccountStorageType::OnChain,
            OFF_CHAIN => AccountStorageType::OffChain,
            _ => panic!("Account with invalid storage bits created"),
        }
    }
    pub fn is_on_chain(&self) -> bool {
        self.storage_type() == AccountStorageType::OnChain
    }
    pub fn get_account_seed(
        init_seed: [u8; 32],
        account_type: AccountType,
        storage_type: AccountStorageType,
        code_commitment: Digest,
        storage_root: Digest,
    ) -> Result<Word, AccountError> {
        get_account_seed(init_seed, account_type, storage_type, code_commitment, storage_root)
    }
    pub fn from_hex(hex_value: &str) -> Result<AccountId, AccountError> {
        hex_to_bytes(hex_value)
            .map_err(|err| AccountError::HexParseError(err.to_string()))
            .and_then(|mut bytes: [u8; 8]| {
                bytes.reverse();
                bytes.try_into()
            })
    }
    pub fn to_hex(&self) -> String {
        format!("0x{:016x}", self.0.as_int())
    }
    pub(super) fn validate_seed_digest(digest: &Digest) -> Result<(), AccountError> {
        let required_zeros = if is_regular_account(digest[0].as_int()) {
            Self::REGULAR_ACCOUNT_SEED_DIGEST_MIN_TRAILING_ZEROS
        } else {
            Self::FAUCET_SEED_DIGEST_MIN_TRAILING_ZEROS
        };
        let trailing_zeros = digest_pow(*digest);
        if required_zeros > trailing_zeros {
            return Err(AccountError::SeedDigestTooFewTrailingZeros {
                expected: required_zeros,
                actual: trailing_zeros,
            });
        }
        Ok(())
    }
}
impl PartialOrd for AccountId {
    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
        Some(self.cmp(other))
    }
}
impl Ord for AccountId {
    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
        self.0.as_int().cmp(&other.0.as_int())
    }
}
impl fmt::Display for AccountId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "0x{:016x}", self.0.as_int())
    }
}
impl From<AccountId> for Felt {
    fn from(id: AccountId) -> Self {
        id.0
    }
}
impl From<AccountId> for [u8; 8] {
    fn from(id: AccountId) -> Self {
        let mut result = [0_u8; 8];
        result[..8].copy_from_slice(&id.0.as_int().to_le_bytes());
        result
    }
}
impl From<AccountId> for u64 {
    fn from(id: AccountId) -> Self {
        id.0.as_int()
    }
}
impl From<AccountId> for LeafIndex<ACCOUNT_TREE_DEPTH> {
    fn from(id: AccountId) -> Self {
        LeafIndex::new_max_depth(id.0.as_int())
    }
}
pub const fn account_id_from_felt(value: Felt) -> Result<AccountId, AccountError> {
    let int_value = value.as_int();
    let count = int_value.count_ones();
    if count < AccountId::MIN_ACCOUNT_ONES {
        return Err(AccountError::AccountIdTooFewOnes(AccountId::MIN_ACCOUNT_ONES, count));
    }
    let bits = (int_value & ACCOUNT_STORAGE_MASK) >> ACCOUNT_STORAGE_MASK_SHIFT;
    match bits {
        ON_CHAIN | OFF_CHAIN => (),
        _ => return Err(AccountError::InvalidAccountStorageType),
    };
    Ok(AccountId(value))
}
impl TryFrom<Felt> for AccountId {
    type Error = AccountError;
    fn try_from(value: Felt) -> Result<Self, Self::Error> {
        account_id_from_felt(value)
    }
}
impl TryFrom<[u8; 8]> for AccountId {
    type Error = AccountError;
    fn try_from(value: [u8; 8]) -> Result<Self, Self::Error> {
        let element = parse_felt(&value[..8])?;
        Self::try_from(element)
    }
}
impl TryFrom<u64> for AccountId {
    type Error = AccountError;
    fn try_from(value: u64) -> Result<Self, Self::Error> {
        let element = parse_felt(&value.to_le_bytes())?;
        Self::try_from(element)
    }
}
impl Serializable for AccountId {
    fn write_into<W: miden_crypto::utils::ByteWriter>(&self, target: &mut W) {
        self.0.write_into(target);
    }
}
impl Deserializable for AccountId {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        Felt::read_from(source)?
            .try_into()
            .map_err(|err: AccountError| DeserializationError::InvalidValue(err.to_string()))
    }
}
fn parse_felt(bytes: &[u8]) -> Result<Felt, AccountError> {
    Felt::try_from(bytes).map_err(|err| AccountError::AccountIdInvalidFieldElement(err.to_string()))
}
pub(super) fn compute_digest(seed: Word, code_commitment: Digest, storage_root: Digest) -> Digest {
    let mut elements = Vec::with_capacity(16);
    elements.extend(seed);
    elements.extend(*code_commitment);
    elements.extend(*storage_root);
    elements.resize(16, ZERO);
    Hasher::hash_elements(&elements)
}
pub(super) fn digest_pow(digest: Digest) -> u32 {
    digest.as_elements()[3].as_int().trailing_zeros()
}
fn is_regular_account(account_id: u64) -> bool {
    let account_type = account_id.into();
    matches!(
        account_type,
        AccountType::RegularAccountUpdatableCode | AccountType::RegularAccountImmutableCode
    )
}
#[cfg(any(feature = "testing", test))]
pub mod testing {
    use super::{
        AccountStorageType, AccountType, ACCOUNT_STORAGE_MASK_SHIFT, ACCOUNT_TYPE_MASK_SHIFT,
    };
    pub const ACCOUNT_ID_SENDER: u64 = account_id(
        AccountType::RegularAccountImmutableCode,
        AccountStorageType::OffChain,
        0b0001_1111,
    );
    pub const ACCOUNT_ID_OFF_CHAIN_SENDER: u64 = account_id(
        AccountType::RegularAccountImmutableCode,
        AccountStorageType::OffChain,
        0b0010_1111,
    );
    pub const ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_OFF_CHAIN: u64 = account_id(
        AccountType::RegularAccountUpdatableCode,
        AccountStorageType::OffChain,
        0b0011_1111,
    );
    pub const ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN: u64 = account_id(
        AccountType::RegularAccountImmutableCode,
        AccountStorageType::OnChain,
        0b0001_1111,
    );
    pub const ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN_2: u64 = account_id(
        AccountType::RegularAccountImmutableCode,
        AccountStorageType::OnChain,
        0b0010_1111,
    );
    pub const ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_ON_CHAIN: u64 = account_id(
        AccountType::RegularAccountUpdatableCode,
        AccountStorageType::OnChain,
        0b0011_1111,
    );
    pub const ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_ON_CHAIN_2: u64 = account_id(
        AccountType::RegularAccountUpdatableCode,
        AccountStorageType::OnChain,
        0b0100_1111,
    );
    pub const ACCOUNT_ID_FUNGIBLE_FAUCET_OFF_CHAIN: u64 =
        account_id(AccountType::FungibleFaucet, AccountStorageType::OffChain, 0b0001_1111);
    pub const ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN: u64 =
        account_id(AccountType::FungibleFaucet, AccountStorageType::OnChain, 0b0001_1111);
    pub const ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN_1: u64 =
        account_id(AccountType::FungibleFaucet, AccountStorageType::OnChain, 0b0010_1111);
    pub const ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN_2: u64 =
        account_id(AccountType::FungibleFaucet, AccountStorageType::OnChain, 0b0011_1111);
    pub const ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN_3: u64 =
        account_id(AccountType::FungibleFaucet, AccountStorageType::OnChain, 0b0100_1111);
    pub const ACCOUNT_ID_INSUFFICIENT_ONES: u64 =
        account_id(AccountType::NonFungibleFaucet, AccountStorageType::OffChain, 0b0000_0000); pub const ACCOUNT_ID_NON_FUNGIBLE_FAUCET_OFF_CHAIN: u64 =
        account_id(AccountType::NonFungibleFaucet, AccountStorageType::OffChain, 0b0001_1111);
    pub const ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN: u64 =
        account_id(AccountType::NonFungibleFaucet, AccountStorageType::OnChain, 0b0010_1111);
    pub const ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN_1: u64 =
        account_id(AccountType::NonFungibleFaucet, AccountStorageType::OnChain, 0b0011_1111);
    pub const fn account_id(
        account_type: AccountType,
        storage: AccountStorageType,
        rest: u64,
    ) -> u64 {
        let mut id = 0;
        id ^= (storage as u64) << ACCOUNT_STORAGE_MASK_SHIFT;
        id ^= (account_type as u64) << ACCOUNT_TYPE_MASK_SHIFT;
        id ^= rest;
        id
    }
}
#[cfg(test)]
mod tests {
    use miden_crypto::utils::{Deserializable, Serializable};
    use super::{
        testing::*, AccountId, AccountStorageType, AccountType, ACCOUNT_ISFAUCET_MASK,
        ACCOUNT_TYPE_MASK_SHIFT, FUNGIBLE_FAUCET, NON_FUNGIBLE_FAUCET,
        REGULAR_ACCOUNT_IMMUTABLE_CODE, REGULAR_ACCOUNT_UPDATABLE_CODE,
    };
    #[test]
    fn test_account_id() {
        use crate::accounts::AccountId;
        for account_type in [
            AccountType::RegularAccountImmutableCode,
            AccountType::RegularAccountUpdatableCode,
            AccountType::NonFungibleFaucet,
            AccountType::FungibleFaucet,
        ] {
            for storage_type in [AccountStorageType::OnChain, AccountStorageType::OffChain] {
                let acc = AccountId::try_from(account_id(account_type, storage_type, 0b1111_1111))
                    .unwrap();
                assert_eq!(acc.account_type(), account_type);
                assert_eq!(acc.storage_type(), storage_type);
            }
        }
    }
    #[test]
    fn test_account_id_from_hex_and_back() {
        for account_id in [
            ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN,
            ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_OFF_CHAIN,
            ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN,
        ] {
            let acc = AccountId::try_from(account_id).expect("Valid account ID");
            assert_eq!(acc, AccountId::from_hex(&acc.to_hex()).unwrap());
        }
    }
    #[test]
    fn test_account_id_serde() {
        let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN)
            .expect("Valid account ID");
        assert_eq!(account_id, AccountId::read_from_bytes(&account_id.to_bytes()).unwrap());
        let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_OFF_CHAIN)
            .expect("Valid account ID");
        assert_eq!(account_id, AccountId::read_from_bytes(&account_id.to_bytes()).unwrap());
        let account_id =
            AccountId::try_from(ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN).expect("Valid account ID");
        assert_eq!(account_id, AccountId::read_from_bytes(&account_id.to_bytes()).unwrap());
        let account_id = AccountId::try_from(ACCOUNT_ID_NON_FUNGIBLE_FAUCET_OFF_CHAIN)
            .expect("Valid account ID");
        assert_eq!(account_id, AccountId::read_from_bytes(&account_id.to_bytes()).unwrap());
    }
    #[test]
    fn test_account_id_account_type() {
        let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN)
            .expect("Valid account ID");
        let account_type: AccountType = ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN.into();
        assert_eq!(account_type, account_id.account_type());
        let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_OFF_CHAIN)
            .expect("Valid account ID");
        let account_type: AccountType = ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_OFF_CHAIN.into();
        assert_eq!(account_type, account_id.account_type());
        let account_id =
            AccountId::try_from(ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN).expect("Valid account ID");
        let account_type: AccountType = ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN.into();
        assert_eq!(account_type, account_id.account_type());
        let account_id = AccountId::try_from(ACCOUNT_ID_NON_FUNGIBLE_FAUCET_OFF_CHAIN)
            .expect("Valid account ID");
        let account_type: AccountType = ACCOUNT_ID_NON_FUNGIBLE_FAUCET_OFF_CHAIN.into();
        assert_eq!(account_type, account_id.account_type());
    }
    #[test]
    fn test_account_id_tag_identifiers() {
        let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN)
            .expect("Valid account ID");
        assert!(account_id.is_regular_account());
        assert_eq!(account_id.account_type(), AccountType::RegularAccountImmutableCode);
        assert!(account_id.is_on_chain());
        let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_OFF_CHAIN)
            .expect("Valid account ID");
        assert!(account_id.is_regular_account());
        assert_eq!(account_id.account_type(), AccountType::RegularAccountUpdatableCode);
        assert!(!account_id.is_on_chain());
        let account_id =
            AccountId::try_from(ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN).expect("Valid account ID");
        assert!(account_id.is_faucet());
        assert_eq!(account_id.account_type(), AccountType::FungibleFaucet);
        assert!(account_id.is_on_chain());
        let account_id = AccountId::try_from(ACCOUNT_ID_NON_FUNGIBLE_FAUCET_OFF_CHAIN)
            .expect("Valid account ID");
        assert!(account_id.is_faucet());
        assert_eq!(account_id.account_type(), AccountType::NonFungibleFaucet);
        assert!(!account_id.is_on_chain());
    }
    #[test]
    fn test_account_id_faucet_bit() {
        assert_ne!((FUNGIBLE_FAUCET << ACCOUNT_TYPE_MASK_SHIFT) & ACCOUNT_ISFAUCET_MASK, 0);
        assert_ne!((NON_FUNGIBLE_FAUCET << ACCOUNT_TYPE_MASK_SHIFT) & ACCOUNT_ISFAUCET_MASK, 0);
        assert_eq!(
            (REGULAR_ACCOUNT_IMMUTABLE_CODE << ACCOUNT_TYPE_MASK_SHIFT) & ACCOUNT_ISFAUCET_MASK,
            0
        );
        assert_eq!(
            (REGULAR_ACCOUNT_UPDATABLE_CODE << ACCOUNT_TYPE_MASK_SHIFT) & ACCOUNT_ISFAUCET_MASK,
            0
        );
    }
}