use std::collections::HashMap;

use near_sdk::{
    borsh::{BorshDeserialize, BorshSerialize},
    collections::LookupMap,
    env,
    json_types::U128,
    near_bindgen, AccountId, BorshStorageKey, Gas, PanicOnDefault, Promise, PromiseResult,
};

type TokenId = String;
type Approval = Option<u64>;

const GAS_FOR_MT_ON_TRANSFER: Gas = Gas::from_tgas(50);
const GAS_FOR_RESOLVE_TRANSFER: Gas = Gas::from_tgas(50);

#[derive(BorshSerialize, BorshStorageKey)]
enum StorageKey {
    TokenBalances,
}

#[near_bindgen]
#[derive(BorshDeserialize, BorshSerialize, PanicOnDefault)]
pub struct MockMultiToken {
    /// Maps from (account_id, token_id) to balance
    balances: LookupMap<String, U128>,
    /// Token supplies by token_id
    supplies: HashMap<String, U128>,
    /// Simulate transfer failures (for testing error scenarios)
    simulate_failure_on_transfer: bool,
}

#[near_bindgen]
impl MockMultiToken {
    #[init]
    pub fn new() -> Self {
        Self {
            balances: LookupMap::new(StorageKey::TokenBalances),
            supplies: HashMap::new(),
            simulate_failure_on_transfer: false,
        }
    }

    // Mock-specific functions for testing
    pub fn mint(&mut self, account_id: AccountId, token_id: TokenId, amount: U128) {
        let key = format!("{}:{}", account_id, token_id);
        let current_balance = self.balances.get(&key).unwrap_or(U128(0));
        self.balances
            .insert(&key, &U128(current_balance.0 + amount.0));

        let current_supply = self.supplies.get(&token_id).unwrap_or(&U128(0));
        self.supplies
            .insert(token_id, U128(current_supply.0 + amount.0));
    }

    pub fn burn(&mut self, account_id: AccountId, token_id: TokenId, amount: U128) {
        let key = format!("{}:{}", account_id, token_id);
        let current_balance = self.balances.get(&key).unwrap_or(U128(0));
        assert!(
            current_balance.0 >= amount.0,
            "Insufficient balance to burn"
        );

        self.balances
            .insert(&key, &U128(current_balance.0 - amount.0));

        let current_supply = self.supplies.get(&token_id).unwrap_or(&U128(0));
        self.supplies
            .insert(token_id, U128(current_supply.0 - amount.0));
    }

    /// Test utility: Enable/disable transfer failure simulation
    pub fn set_simulate_transfer_failure(&mut self, should_fail: bool) {
        self.simulate_failure_on_transfer = should_fail;
    }

    // NEP-245 Core View Methods
    pub fn mt_token(&self, token_ids: Vec<TokenId>) -> Vec<bool> {
        token_ids
            .into_iter()
            .map(|token_id| self.supplies.contains_key(&token_id))
            .collect()
    }

    pub fn mt_balance_of(&self, account_id: AccountId, token_id: TokenId) -> U128 {
        let key = format!("{}:{}", account_id, token_id);
        self.balances.get(&key).unwrap_or(U128(0))
    }

    pub fn mt_batch_balance_of(&self, account_id: AccountId, token_ids: Vec<TokenId>) -> Vec<U128> {
        token_ids
            .into_iter()
            .map(|token_id| self.mt_balance_of(account_id.clone(), token_id))
            .collect()
    }

    pub fn mt_supply(&self, token_id: TokenId) -> Option<U128> {
        self.supplies.get(&token_id).cloned()
    }

    pub fn mt_batch_supply(&self, token_ids: Vec<TokenId>) -> Vec<Option<U128>> {
        token_ids
            .into_iter()
            .map(|token_id| self.mt_supply(token_id))
            .collect()
    }

    // NEP-245 Core Change Methods
    #[payable]
    pub fn mt_transfer(
        &mut self,
        receiver_id: AccountId,
        token_id: TokenId,
        amount: U128,
        _approval: Approval,
        _memo: Option<String>,
    ) {
        // Require 1 yoctoⓃ deposit for security
        assert_eq!(
            env::attached_deposit().as_yoctonear(),
            1,
            "Requires attached deposit of exactly 1 yoctoⓃ"
        );

        let sender = env::predecessor_account_id();
        self.internal_transfer(&sender, &receiver_id, &token_id, amount);
    }

    #[payable]
    pub fn mt_batch_transfer(
        &mut self,
        receiver_id: AccountId,
        token_ids: Vec<TokenId>,
        amounts: Vec<U128>,
        _approvals: Option<Vec<Approval>>,
        _memo: Option<String>,
    ) {
        // Require 1 yoctoⓃ deposit for security
        assert_eq!(
            env::attached_deposit().as_yoctonear(),
            1,
            "Requires attached deposit of exactly 1 yoctoⓃ"
        );

        assert_eq!(
            token_ids.len(),
            amounts.len(),
            "Token IDs and amounts must have the same length"
        );

        let sender = env::predecessor_account_id();
        for (token_id, amount) in token_ids.iter().zip(amounts.iter()) {
            self.internal_transfer(&sender, &receiver_id, token_id, *amount);
        }
    }

    #[payable]
    pub fn mt_transfer_call(
        &mut self,
        receiver_id: AccountId,
        token_id: TokenId,
        amount: U128,
        approval: Approval,
        memo: Option<String>,
        msg: String,
    ) -> Promise {
        let sender = env::predecessor_account_id();
        self.mt_transfer(
            receiver_id.clone(),
            token_id.clone(),
            amount,
            approval,
            memo,
        );

        // Call the receiver contract's mt_on_transfer method
        Promise::new(receiver_id.clone())
            .function_call(
                "mt_on_transfer".to_string(),
                near_sdk::serde_json::to_vec(&near_sdk::serde_json::json!({
                    "sender_id": sender,
                    "previous_owner_ids": vec![sender.clone()],
                    "token_ids": vec![token_id.clone()],
                    "amounts": vec![amount],
                    "msg": msg
                }))
                .unwrap(),
                near_sdk::NearToken::from_yoctonear(0),
                GAS_FOR_MT_ON_TRANSFER,
            )
            .then(
                Promise::new(env::current_account_id()).function_call(
                    "mt_resolve_transfer".to_string(),
                    near_sdk::serde_json::to_vec(&near_sdk::serde_json::json!({
                        "sender_id": sender,
                        "receiver_id": receiver_id,
                        "token_ids": vec![token_id],
                        "amounts": vec![amount]
                    }))
                    .unwrap(),
                    near_sdk::NearToken::from_yoctonear(0),
                    GAS_FOR_RESOLVE_TRANSFER,
                ),
            )
    }

    #[payable]
    pub fn mt_batch_transfer_call(
        &mut self,
        receiver_id: AccountId,
        token_ids: Vec<TokenId>,
        amounts: Vec<U128>,
        approvals: Option<Vec<Approval>>,
        memo: Option<String>,
        msg: String,
    ) -> Promise {
        let sender = env::predecessor_account_id();
        self.mt_batch_transfer(
            receiver_id.clone(),
            token_ids.clone(),
            amounts.clone(),
            approvals,
            memo,
        );

        // Call the receiver contract's mt_on_transfer method
        Promise::new(receiver_id.clone())
            .function_call(
                "mt_on_transfer".to_string(),
                near_sdk::serde_json::to_vec(&near_sdk::serde_json::json!({
                    "sender_id": sender,
                    "previous_owner_ids": vec![sender.clone()],
                    "token_ids": token_ids.clone(),
                    "amounts": amounts.clone(),
                    "msg": msg
                }))
                .unwrap(),
                near_sdk::NearToken::from_yoctonear(0),
                GAS_FOR_MT_ON_TRANSFER,
            )
            .then(
                Promise::new(env::current_account_id()).function_call(
                    "mt_resolve_batch_transfer".to_string(),
                    near_sdk::serde_json::to_vec(&near_sdk::serde_json::json!({
                        "sender_id": sender,
                        "receiver_id": receiver_id,
                        "token_ids": token_ids,
                        "amounts": amounts
                    }))
                    .unwrap(),
                    near_sdk::NearToken::from_yoctonear(0),
                    GAS_FOR_RESOLVE_TRANSFER,
                ),
            )
    }

    // Promise resolution callbacks
    #[private]
    pub fn mt_resolve_transfer(
        &mut self,
        sender_id: AccountId,
        receiver_id: AccountId,
        token_ids: Vec<TokenId>,
        amounts: Vec<U128>,
    ) -> Vec<U128> {
        match env::promise_result(0) {
            PromiseResult::Successful(result) => {
                // Try to parse the result as unused amounts Vec<U128>
                if let Ok(unused_amounts) = near_sdk::serde_json::from_slice::<Vec<U128>>(&result) {
                    // Refund unused tokens
                    for (i, (token_id, _amount)) in token_ids.iter().zip(amounts.iter()).enumerate()
                    {
                        if let Some(unused) = unused_amounts.get(i) {
                            if unused.0 > 0 {
                                self.internal_transfer(&receiver_id, &sender_id, token_id, *unused);
                            }
                        }
                    }

                    // Return amounts used (total - unused)
                    amounts
                        .iter()
                        .zip(unused_amounts.iter())
                        .map(|(total, unused)| U128(total.0 - unused.0))
                        .collect()
                } else {
                    // No unused amounts vector returned, assume all was used
                    amounts
                }
            }
            PromiseResult::Failed => {
                // Transfer failed, refund all tokens
                for (token_id, amount) in token_ids.iter().zip(amounts.iter()) {
                    self.internal_transfer(&receiver_id, &sender_id, token_id, *amount);
                }

                // Transfer failed, so nothing was used
                vec![U128(0); token_ids.len()]
            }
        }
    }

    #[private]
    pub fn mt_resolve_batch_transfer(
        &mut self,
        sender_id: AccountId,
        receiver_id: AccountId,
        token_ids: Vec<TokenId>,
        amounts: Vec<U128>,
    ) -> Vec<U128> {
        // Batch resolve uses same logic as single resolve
        self.mt_resolve_transfer(sender_id, receiver_id, token_ids, amounts)
    }

    // Internal helper methods
    fn internal_transfer(
        &mut self,
        sender_id: &AccountId,
        receiver_id: &AccountId,
        token_id: &TokenId,
        amount: U128,
    ) {
        // Simulate transfer failure for testing
        if self.simulate_failure_on_transfer {
            panic!("Simulated transfer failure");
        }

        // Validate token exists
        assert!(self.supplies.contains_key(token_id), "Token does not exist");

        // Simulate storage deposit requirement - fail if receiver is "nonexistent.testnet"
        assert!(
            receiver_id.as_str() != "nonexistent.testnet",
            "Account does not exist or has no storage deposit"
        );

        let sender_key = format!("{}:{}", sender_id, token_id);
        let receiver_key = format!("{}:{}", receiver_id, token_id);

        let sender_balance = self.balances.get(&sender_key).unwrap_or(U128(0));
        assert!(sender_balance.0 >= amount.0, "Insufficient balance");

        // Perform the transfer
        self.balances
            .insert(&sender_key, &U128(sender_balance.0 - amount.0));

        let receiver_balance = self.balances.get(&receiver_key).unwrap_or(U128(0));
        self.balances
            .insert(&receiver_key, &U128(receiver_balance.0 + amount.0));
    }
}

// Tests
#[cfg(test)]
mod tests {
    use super::*;
    use near_sdk::test_utils::{accounts, VMContextBuilder};
    use near_sdk::{testing_env, NearToken};

    fn get_context(predecessor_account_id: AccountId) -> VMContextBuilder {
        let mut builder = VMContextBuilder::new();
        builder
            .current_account_id(accounts(0))
            .signer_account_id(predecessor_account_id.clone())
            .predecessor_account_id(predecessor_account_id);
        builder
    }

    #[test]
    fn test_mint_and_balance() {
        let context = get_context(accounts(1));
        testing_env!(context.build());

        let mut contract = MockMultiToken::new();

        // Test minting
        contract.mint(accounts(1), "token1".to_string(), U128(1000));

        // Test balance query
        let balance = contract.mt_balance_of(accounts(1), "token1".to_string());
        assert_eq!(balance.0, 1000);

        // Test supply query
        let supply = contract.mt_supply("token1".to_string()).unwrap();
        assert_eq!(supply.0, 1000);

        // Test token existence
        let exists = contract.mt_token(vec!["token1".to_string()]);
        assert_eq!(exists, vec![true]);
    }

    #[test]
    fn test_transfer() {
        let mut context = get_context(accounts(1));
        context.attached_deposit(NearToken::from_yoctonear(1));
        testing_env!(context.build());

        let mut contract = MockMultiToken::new();

        // Mint tokens to sender
        contract.mint(accounts(1), "token1".to_string(), U128(1000));

        // Transfer tokens
        contract.mt_transfer(accounts(2), "token1".to_string(), U128(500), None, None);

        // Check balances
        let sender_balance = contract.mt_balance_of(accounts(1), "token1".to_string());
        let receiver_balance = contract.mt_balance_of(accounts(2), "token1".to_string());

        assert_eq!(sender_balance.0, 500);
        assert_eq!(receiver_balance.0, 500);
    }

    #[test]
    fn test_batch_operations() {
        let mut context = get_context(accounts(1));
        context.attached_deposit(NearToken::from_yoctonear(1));
        testing_env!(context.build());

        let mut contract = MockMultiToken::new();

        // Mint multiple tokens
        contract.mint(accounts(1), "token1".to_string(), U128(1000));
        contract.mint(accounts(1), "token2".to_string(), U128(2000));

        // Test batch balance query
        let balances = contract.mt_batch_balance_of(
            accounts(1),
            vec!["token1".to_string(), "token2".to_string()],
        );
        assert_eq!(balances, vec![U128(1000), U128(2000)]);

        // Test batch transfer
        contract.mt_batch_transfer(
            accounts(2),
            vec!["token1".to_string(), "token2".to_string()],
            vec![U128(300), U128(700)],
            None,
            None,
        );

        // Check final balances
        let final_balances = contract.mt_batch_balance_of(
            accounts(2),
            vec!["token1".to_string(), "token2".to_string()],
        );
        assert_eq!(final_balances, vec![U128(300), U128(700)]);
    }
}