use near_sdk::{
    AccountId, Gas, IntoStorageKey, PromiseOrValue, StorageUsage, assert_one_yocto,
    collections::LookupMap, env, json_types::U128, log, near, require, serde_json,
};

use crate::fungible_token::{
    core::FungibleTokenCore,
    events::{FtBurn, FtTransfer},
    receiver::ext_ft_receiver,
    resolver::{FungibleTokenResolver, ext_ft_resolver},
};

const GAS_FOR_RESOLVE_TRANSFER: Gas = Gas::from_tgas(5);

const ERR_TOTAL_SUPPLY_OVERFLOW: &str = "Total supply overflow";

pub type Balance = u128;

/// Implementation of a FungibleToken standard.
/// Allows to include NEP-141 compatible token to any contract.
/// There are next traits that any contract may implement:
///     - FungibleTokenCore -- interface with ft_transfer methods. FungibleToken provides methods for it.
///     - FungibleTokenMetaData -- return metadata for the token in NEP-148, up to contract to implement.
///     - StorageManager -- interface for NEP-145 for allocating storage per account. FungibleToken provides methods for it.
///     - AccountRegistrar -- interface for an account to register and unregister
///
/// For example usage, see examples/fungible-token/src/lib.rs.
#[near]
pub struct FungibleToken {
    /// sha256(AccountId) -> Account balance
    pub accounts: LookupMapAdapter,

    /// Total supply of the all token.
    pub total_supply: Balance,

    /// The storage size in bytes for one account.
    pub account_storage_usage: StorageUsage,
}

/// Storage key wrapper for an account balance.
///
/// Kept as an enum so the stored key retains its leading borsh discriminant
/// byte (`Hash` => `0x00`). Existing balances are keyed by `0x00 ‖ sha256` on
/// the trie; collapsing this into a bare `[u8; 32]` would drop that byte and
/// orphan every entry, which can't be re-keyed cheaply (the map isn't
/// enumerable). The historical `AccountId` variant has been removed — no
/// account was ever stored unhashed — and removing it leaves `Hash` at
/// discriminant `0`, so the encoding is unchanged.
#[near]
pub enum LookupMapKey {
    Hash([u8; 32]),
}

#[near]
pub struct LookupMapAdapter {
    inner: LookupMap<LookupMapKey, Balance>,
}

impl FungibleToken {
    pub fn new<S>(prefix: S) -> Self
    where
        S: IntoStorageKey,
    {
        let mut this = Self {
            accounts: LookupMapAdapter::new(prefix),
            total_supply: 0,
            account_storage_usage: 0,
        };
        this.measure_account_storage_usage();
        this
    }

    /// Reconstructs the token over already-populated `accounts` storage during a
    /// state migration. Unlike [`FungibleToken::new`], it adopts `prefix`
    /// without re-measuring per-account storage — measurement transiently writes
    /// a probe key and would clobber a live account on a populated map — so the
    /// recorded `total_supply` and `account_storage_usage` are carried over from
    /// the previous state instead.
    pub fn from_prefix<S>(
        prefix: S,
        total_supply: Balance,
        account_storage_usage: StorageUsage,
    ) -> Self
    where
        S: IntoStorageKey,
    {
        Self { accounts: LookupMapAdapter::new(prefix), total_supply, account_storage_usage }
    }

    fn measure_account_storage_usage(&mut self) {
        let initial_storage_usage = env::storage_usage();
        let tmp_account_id = "a".repeat(64).parse().unwrap();
        self.accounts.insert(&tmp_account_id, &0u128);
        self.account_storage_usage = env::storage_usage() - initial_storage_usage;
        self.accounts.remove(&tmp_account_id);
    }

    pub fn internal_unwrap_balance_of(&self, account_id: &AccountId) -> Balance {
        match self.accounts.get(account_id) {
            Some(balance) => balance,
            None => {
                env::panic_str(format!("The account {} is not registered", &account_id).as_str())
            }
        }
    }

    pub fn internal_deposit(&mut self, account_id: &AccountId, amount: Balance) {
        let balance = self.internal_unwrap_balance_of(account_id);
        if let Some(new_balance) = balance.checked_add(amount) {
            self.accounts.insert(account_id, &new_balance);
            self.total_supply = self
                .total_supply
                .checked_add(amount)
                .unwrap_or_else(|| env::panic_str(ERR_TOTAL_SUPPLY_OVERFLOW));
        } else {
            env::panic_str("Balance overflow");
        }
    }

    pub fn internal_withdraw(&mut self, account_id: &AccountId, amount: Balance) {
        let balance = self.internal_unwrap_balance_of(account_id);
        if let Some(new_balance) = balance.checked_sub(amount) {
            self.accounts.insert(account_id, &new_balance);
            self.total_supply = self
                .total_supply
                .checked_sub(amount)
                .unwrap_or_else(|| env::panic_str(ERR_TOTAL_SUPPLY_OVERFLOW));
        } else {
            env::panic_str("The account doesn't have enough balance");
        }
    }

    pub fn internal_transfer(
        &mut self,
        sender_id: &AccountId,
        receiver_id: &AccountId,
        amount: Balance,
        memo: Option<String>,
    ) {
        require!(sender_id != receiver_id, "Sender and receiver should be different");
        require!(amount > 0, "The amount should be a positive number");
        self.internal_withdraw(sender_id, amount);
        self.internal_deposit(receiver_id, amount);
        FtTransfer {
            old_owner_id: sender_id,
            new_owner_id: receiver_id,
            amount: U128(amount),
            memo: memo.as_deref(),
        }
        .emit();
    }

    pub fn internal_register_account(&mut self, account_id: &AccountId) {
        if self.accounts.insert(account_id, &0).is_some() {
            env::panic_str("The account is already registered");
        }
    }
}

impl FungibleTokenCore for FungibleToken {
    fn ft_transfer(&mut self, receiver_id: AccountId, amount: U128, memo: Option<String>) {
        assert_one_yocto();
        let sender_id = env::predecessor_account_id();
        let amount: Balance = amount.into();
        self.internal_transfer(&sender_id, &receiver_id, amount, memo);
    }

    fn ft_transfer_call(
        &mut self,
        receiver_id: AccountId,
        amount: U128,
        memo: Option<String>,
        msg: String,
    ) -> PromiseOrValue<U128> {
        assert_one_yocto();

        let sender_id = env::predecessor_account_id();
        let amount: Balance = amount.into();
        self.internal_transfer(&sender_id, &receiver_id, amount, memo);
        // Initiating receiver's call and the callback
        ext_ft_receiver::ext(receiver_id.clone())
            // forward all remaining gas to `ft_on_transfer`
            .with_unused_gas_weight(1)
            .ft_on_transfer(sender_id.clone(), amount.into(), msg)
            .then(
                ext_ft_resolver::ext(env::current_account_id())
                    .with_static_gas(GAS_FOR_RESOLVE_TRANSFER)
                    // do not distribute remaining gas for `ft_resolve_transfer`
                    .with_unused_gas_weight(0)
                    .ft_resolve_transfer(sender_id, receiver_id, amount.into()),
            )
            .into()
    }

    fn ft_total_supply(&self) -> U128 {
        self.total_supply.into()
    }

    fn ft_balance_of(&self, account_id: AccountId) -> U128 {
        self.accounts.get(&account_id).unwrap_or(0).into()
    }
}

impl FungibleToken {
    /// Internal method that returns the amount of burned tokens in a corner case when the sender
    /// has deleted (unregistered) their account while the `ft_transfer_call` was still in flight.
    /// Returns (Used token amount, Burned token amount)
    pub fn internal_ft_resolve_transfer(
        &mut self,
        sender_id: &AccountId,
        receiver_id: AccountId,
        amount: U128,
    ) -> (u128, u128) {
        const MAX_RESULT_LENGTH: usize = "\"+340282366920938463463374607431768211455\"".len(); // u128::MAX

        // Get the unused amount from the `ft_on_transfer` call result.
        let unused_amount = env::promise_result_checked(0, MAX_RESULT_LENGTH)
            .ok()
            .and_then(|value| serde_json::from_slice::<U128>(&value).ok())
            .unwrap_or(amount)
            .0
            .min(amount.0);

        let amount: Balance = amount.into();

        if unused_amount > 0 {
            let receiver_balance = self.accounts.get(&receiver_id).unwrap_or(0);
            if receiver_balance > 0 {
                let refund_amount = std::cmp::min(receiver_balance, unused_amount);
                if let Some(new_receiver_balance) = receiver_balance.checked_sub(refund_amount) {
                    self.accounts.insert(&receiver_id, &new_receiver_balance);
                } else {
                    env::panic_str("The receiver account doesn't have enough balance");
                }

                if let Some(sender_balance) = self.accounts.get(sender_id) {
                    if let Some(new_sender_balance) = sender_balance.checked_add(refund_amount) {
                        self.accounts.insert(sender_id, &new_sender_balance);
                    } else {
                        env::panic_str("Sender balance overflow");
                    }

                    FtTransfer {
                        old_owner_id: &receiver_id,
                        new_owner_id: sender_id,
                        amount: U128(refund_amount),
                        memo: Some("refund"),
                    }
                    .emit();
                    let used_amount = amount
                        .checked_sub(refund_amount)
                        .unwrap_or_else(|| env::panic_str(ERR_TOTAL_SUPPLY_OVERFLOW));
                    return (used_amount, 0);
                } else {
                    // Sender's account was deleted, so we need to burn tokens.
                    self.total_supply = self
                        .total_supply
                        .checked_sub(refund_amount)
                        .unwrap_or_else(|| env::panic_str(ERR_TOTAL_SUPPLY_OVERFLOW));
                    log!("The account of the sender was deleted");
                    FtBurn {
                        owner_id: &receiver_id,
                        amount: U128(refund_amount),
                        memo: Some("refund"),
                    }
                    .emit();
                    return (amount, refund_amount);
                }
            }
        }
        (amount, 0)
    }
}

impl FungibleTokenResolver for FungibleToken {
    fn ft_resolve_transfer(
        &mut self,
        sender_id: AccountId,
        receiver_id: AccountId,
        amount: U128,
    ) -> U128 {
        self.internal_ft_resolve_transfer(&sender_id, receiver_id, amount).0.into()
    }
}

impl LookupMapAdapter {
    fn new<S: IntoStorageKey>(prefix: S) -> LookupMapAdapter {
        Self { inner: LookupMap::new(prefix) }
    }

    fn hash_key(account: &AccountId) -> LookupMapKey {
        LookupMapKey::Hash(env::sha256_array(account.as_bytes()))
    }

    pub fn get(&self, key: &AccountId) -> Option<Balance> {
        self.inner.get(&Self::hash_key(key))
    }

    pub fn remove(&mut self, key: &AccountId) -> Option<Balance> {
        self.inner.remove(&Self::hash_key(key))
    }

    pub fn insert(&mut self, key: &AccountId, value: &Balance) -> Option<Balance> {
        self.inner.insert(&Self::hash_key(key), value)
    }

    /// Returns true if the map contains a given key.
    pub fn contains_key(&self, key: &AccountId) -> bool {
        self.inner.contains_key(&Self::hash_key(key))
    }
}