mod entry;
mod nonces;

pub use self::{entry::*, nonces::MaybeLegacyAccountNonces};

use std::borrow::Cow;

use bitflags::bitflags;
use defuse_bitmap::U256;
use defuse_core::{
    NoncePrefix, Result, accounts::AccountEvent, crypto::PublicKey, events::DefuseEvent,
    intents::account::SetAuthByPredecessorId,
};

use defuse_near_utils::NestPrefix;
use impl_tools::autoimpl;
use near_account_id::AccountType;
use near_sdk::{
    AccountIdRef, BorshStorageKey, IntoStorageKey,
    borsh::BorshSerialize,
    near,
    store::{IterableSet, LookupMap},
};

use super::AccountState;

// NOTE: in order to migrate to a new version (even when adding new fields),
// see docs for `VersionedAccountEntry`
#[derive(Debug)]
#[near(serializers = [borsh])]
#[autoimpl(Deref using self.state)]
#[autoimpl(DerefMut using self.state)]
pub struct Account {
    nonces: MaybeLegacyAccountNonces,

    flags: AccountFlags,
    public_keys: IterableSet<PublicKey>,

    pub state: AccountState,

    prefix: Vec<u8>,
}

impl Account {
    #[inline]
    pub fn new<S>(prefix: S, me: &AccountIdRef) -> Self
    where
        S: IntoStorageKey,
    {
        let prefix = prefix.into_storage_key();

        Self {
            nonces: MaybeLegacyAccountNonces::new(LookupMap::with_hasher(
                prefix.as_slice().nest(AccountPrefix::OptimizedNonces),
            )),
            flags: (!has_implicit_public_key(me))
                .then_some(AccountFlags::IMPLICIT_PUBLIC_KEY_REMOVED)
                .unwrap_or_else(AccountFlags::empty),
            public_keys: IterableSet::new(prefix.as_slice().nest(AccountPrefix::PublicKeys)),
            state: AccountState::new(prefix.as_slice().nest(AccountPrefix::State)),
            prefix,
        }
    }

    #[inline]
    #[must_use]
    pub fn add_public_key(&mut self, me: &AccountIdRef, public_key: PublicKey) -> bool {
        if me == public_key.to_implicit_account_id() {
            let was_removed = self.is_implicit_public_key_removed();
            self.set_implicit_public_key_removed(false);
            was_removed
        } else {
            self.public_keys.insert(public_key)
        }
    }

    #[inline]
    #[must_use]
    pub fn remove_public_key(&mut self, me: &AccountIdRef, public_key: &PublicKey) -> bool {
        if me == public_key.to_implicit_account_id() {
            let was_removed = self.is_implicit_public_key_removed();
            self.set_implicit_public_key_removed(true);
            !was_removed
        } else {
            self.public_keys.remove(public_key)
        }
    }

    #[inline]
    pub fn has_public_key(&self, me: &AccountIdRef, public_key: &PublicKey) -> bool {
        !self.is_implicit_public_key_removed() && me == public_key.to_implicit_account_id()
            || self.public_keys.contains(public_key)
    }

    #[inline]
    pub fn iter_public_keys(&self, me: &AccountIdRef) -> impl Iterator<Item = PublicKey> + '_ {
        self.public_keys.iter().copied().chain(
            (!self.is_implicit_public_key_removed())
                .then(|| PublicKey::from_implicit_account_id(me))
                .flatten(),
        )
    }

    #[inline]
    pub fn is_nonce_used(&self, nonce: U256) -> bool {
        self.nonces.is_used(nonce)
    }

    #[inline]
    pub fn commit_nonce(&mut self, nonce: U256) -> Result<()> {
        self.nonces.commit(nonce)
    }

    /// Clears the all nonces with corresponding prefix if it was expired/invalidated.
    /// Returns whether the nonces was cleared,
    /// regardless of whether it was previously committed or not.
    #[inline]
    pub fn cleanup_nonce_by_prefix(&mut self, prefix: NoncePrefix) -> bool {
        self.nonces.cleanup_by_prefix(prefix)
    }

    #[inline]
    const fn is_implicit_public_key_removed(&self) -> bool {
        self.flags
            .contains(AccountFlags::IMPLICIT_PUBLIC_KEY_REMOVED)
    }

    #[inline]
    fn set_implicit_public_key_removed(&mut self, removed: bool) {
        self.flags
            .set(AccountFlags::IMPLICIT_PUBLIC_KEY_REMOVED, removed);
    }

    /// Returns whether authentication by PREDECESSOR is enabled.
    pub const fn is_auth_by_predecessor_id_enabled(&self) -> bool {
        !self
            .flags
            .contains(AccountFlags::AUTH_BY_PREDECESSOR_ID_DISABLED)
    }

    /// Sets whether authentication by `PREDECESSOR_ID` is enabled.
    /// Returns whether authentication by `PREDECESSOR_ID` was enabled
    /// before.
    pub fn set_auth_by_predecessor_id(&mut self, me: &AccountIdRef, enable: bool) -> bool {
        let was_enabled = self.is_auth_by_predecessor_id_enabled();
        let toggle = was_enabled ^ enable;
        if toggle {
            self.flags
                .toggle(AccountFlags::AUTH_BY_PREDECESSOR_ID_DISABLED);

            DefuseEvent::SetAuthByPredecessorId(AccountEvent::new(
                Cow::Borrowed(me),
                SetAuthByPredecessorId { enabled: enable },
            ))
            .emit();
        }
        was_enabled
    }
}

fn has_implicit_public_key(account: &AccountIdRef) -> bool {
    match account.get_account_type(){
        AccountType::NearImplicitAccount | AccountType::EthImplicitAccount => true,
        AccountType::NamedAccount /* | AccountType::NearDeterministicAccount */ => false
    }
}

#[allow(deprecated)]
mod prefix {
    use super::{BorshSerialize, BorshStorageKey};

    #[derive(BorshSerialize, BorshStorageKey)]
    #[borsh(crate = "::near_sdk::borsh")]
    pub enum AccountPrefix {
        #[deprecated(note = "Please use `AccountPrefix::OptimizedNonces` instead.")]
        _LegacyNonces,
        PublicKeys,
        State,
        OptimizedNonces,
    }
}

use prefix::AccountPrefix;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[near(serializers = [borsh])]
#[repr(transparent)]
struct AccountFlags(u8);

bitflags! {
    impl AccountFlags: u8 {
        // It was a legacy `implicit_public_key_removed: bool`
        // flag in previous version. It's safe to migrate here,
        // since borsh serializes `bool` to 0u8/1u8
        const IMPLICIT_PUBLIC_KEY_REMOVED     = 1 << 0;
        const AUTH_BY_PREDECESSOR_ID_DISABLED = 1 << 1;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use near_sdk::borsh;
    use rstest::rstest;

    #[rstest]
    #[test]
    fn upgrade_to_flags(#[values(true, false)] implicit_public_key_removed: bool) {
        let serialized_legacy = borsh::to_vec(&implicit_public_key_removed).unwrap();
        let flags: AccountFlags = borsh::from_slice(&serialized_legacy).unwrap();
        assert_eq!(
            flags.contains(AccountFlags::IMPLICIT_PUBLIC_KEY_REMOVED),
            implicit_public_key_removed,
            "implicit_public_key_removed doesn't match"
        );
        assert_eq!(
            borsh::to_vec(&flags).unwrap(),
            serialized_legacy,
            "unknown flags set"
        );
    }

    #[test]
    fn near_implicit_account_has_implicit_public_key() {
        assert!(
            !Account::new(
                b"prefix".to_vec(),
                AccountIdRef::new_or_panic(
                    "98793cd91a3f870fb126f66285808c7e094afcfc4eda8a970f6648cdf0dbd6de"
                ),
            )
            .is_implicit_public_key_removed()
        );
    }

    #[test]
    fn eth_implicit_account_has_implicit_public_key() {
        assert!(
            !Account::new(
                b"prefix".to_vec(),
                AccountIdRef::new_or_panic("0x0000000000000000000000000000000000000001"),
            )
            .is_implicit_public_key_removed()
        );
    }

    #[test]
    fn named_account_has_no_implicit_public_key() {
        assert!(
            Account::new(b"prefix".to_vec(), AccountIdRef::new_or_panic("alice.near"),)
                .is_implicit_public_key_removed()
        );
    }
}