use crate::data::{
    address::UserAddressArgs,
    user::{UserData, UserDataArgs},
};
use crate::error::UserStateError;
use std::collections::BTreeMap;

#[cfg(test)]
use crate::mocks::ic_caller;

#[cfg(not(test))]
use ic_cdk::caller as ic_caller;

use ic_cdk::export::{
    candid::{CandidType, Principal},
    serde::Deserialize,
};

/// Represents the state of the canister, including users, and wallet canister.
#[derive(CandidType, Deserialize, Debug)]
pub struct UserState {
    pub users: BTreeMap<Principal, UserData>,
    pub owner: Principal,
    pub wallet_canister: Principal,
}

impl UserState {
    /// Returns the default state.
    pub fn default() -> Self {
        Self {
            users: BTreeMap::default(),
            owner: ic_caller(),
            wallet_canister: Principal::management_canister(),
        }
    }

    /// Initializes the state with the wallet canister.
    /// This function should only be called once, when the canister is first installed.
    pub fn init(&mut self, wallet_canister: Principal) {
        self.wallet_canister = wallet_canister;
    }

    /// Checks if the caller is the wallet canister.
    fn is_caller_wallet_canister(&self) -> bool {
        ic_caller() == self.wallet_canister
    }

    /// Checks if the caller is the owner.
    fn is_caller_owner(&self) -> bool {
        ic_caller() == self.owner
    }

    /// Checks if the caller is the user.
    fn is_caller_user(&self, user: &Principal) -> bool {
        ic_caller() == *user
    }

    /// Validates if the caller is the owner.
    pub fn validate_caller_owner(&self) -> Result<(), UserStateError> {
        if !self.is_caller_owner() {
            return Err(UserStateError::CallerIsNotOwner);
        }

        Ok(())
    }

    /// Validates if the caller is the  wallet canister.
    pub fn validate_caller_wallet_canister(&self) -> Result<(), UserStateError> {
        if !self.is_caller_wallet_canister() {
            return Err(UserStateError::CallerIsNotWalletCanister);
        }

        Ok(())
    }

    /// Validates if the caller is the wallet canister, or the user.
    pub fn validate_caller_wallet_canister_or_user(
        &self,
        user: &Principal,
    ) -> Result<(), UserStateError> {
        if !self.is_caller_wallet_canister() && !self.is_caller_user(user) {
            return Err(UserStateError::CallerNotAuthorized);
        }

        Ok(())
    }

    /// Changes the owner of the canister.    
    /// This function should only be called by the owner.
    pub fn change_owner(&mut self, new_owner: Principal) -> Result<Principal, UserStateError> {
        self.validate_caller_owner()?;

        self.owner = new_owner;

        Ok(new_owner)
    }

    /// Changes the wallet canister.
    /// This function should only be called by the owner.
    pub fn change_wallet_canister(
        &mut self,
        new_wallet_canister: Principal,
    ) -> Result<Principal, UserStateError> {
        self.validate_caller_owner()?;

        self.wallet_canister = new_wallet_canister;

        Ok(new_wallet_canister)
    }

    /// This function returns a reference to the user data.
    /// Adds new user data for a given user Principal.
    pub fn create_user(
        &mut self,
        user: Principal,
        user_args: UserDataArgs,
        address_args: UserAddressArgs,
    ) -> Result<UserData, UserStateError> {
        self.validate_caller_wallet_canister()?;

        if self.users.contains_key(&user) {
            return Err(UserStateError::UserAlreadyExists);
        }

        let user_data = UserData::new(user_args, address_args);
        self.users.insert(user, user_data);

        Ok(self.users.get(&user).unwrap().clone())
    }

    /// This function returns a mutable reference to the user data.
    /// Retrieves user data for a given user Principal.
    pub fn get_user_mut(&mut self, user: &Principal) -> Result<&mut UserData, UserStateError> {
        self.validate_caller_wallet_canister()?;

        self.users.get_mut(user).ok_or(UserStateError::UserNotFound)
    }

    /// This function returns a reference to the user data.
    /// Retrieves user data for a given user Principal.
    pub fn get_user(&self, user: &Principal) -> Result<&UserData, UserStateError> {
        self.validate_caller_wallet_canister_or_user(user)?;

        self.users
            .get(user)
            .map(|user_data| user_data)
            .ok_or(UserStateError::UserNotFound)
    }
}

#[cfg(test)]
mod tests {
    use crate::{mocks::*, state::UserState};

    use super::*;
    use proptest::prelude::*;

    fn initialize_state() -> UserState {
        owner_caller();

        let mut state = UserState::default();
        state.init(wallet_canister_principal());

        state
    }

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(5))]

        #[test]
        fn test_user_state_validate_caller_owner(principal in principal_strategy()) {
            let mut user_state = initialize_state();

            owner_caller();

            assert!(user_state.validate_caller_owner().is_ok());

            random_caller();

            assert!(user_state.validate_caller_owner().is_err());

            owner_caller();

            let new_owner = user_state.change_owner(principal).unwrap();

            assert_eq!(user_state.owner, principal);
            assert_eq!(user_state.owner, new_owner);
        }

        #[test]
        fn test_user_state_validate_caller_wallet_canister(_ in ".*") {
            let mut user_state = initialize_state();

            let principal = wallet_canister_principal();

            random_caller();

            user_state.init(principal);

            assert!(user_state.validate_caller_wallet_canister().is_err());

            wallet_canister_caller();

            assert!(user_state.validate_caller_wallet_canister().is_ok());

            let random_data = UserDataArgs::default();

            let randdom_address = UserAddressArgs::default();

            assert!(user_state.create_user(
                principal,
                random_data,
                randdom_address,
            ).is_ok());

            random_caller();

            assert!(user_state.get_user_mut(&principal).is_err());
        }

        #[test]
        fn test_user_state_validate_caller_wallet_canister_or_user(principal in principal_strategy()) {
            let mut user_state = initialize_state();

            set_caller(principal);

            assert!(user_state.validate_caller_wallet_canister_or_user(&principal).is_ok());

            let random_data = UserDataArgs::default();

            let randdom_address = UserAddressArgs::default();

            wallet_canister_caller();

            assert!(user_state.create_user(
                principal,
                random_data,
                randdom_address,
            ).is_ok());

            assert!(user_state.get_user(&principal).is_ok());

            random_caller();

            assert!(user_state.get_user(&principal).is_err());

            owner_caller();

            assert!(user_state.get_user(&principal).is_err());

            wallet_canister_caller();

            assert!(user_state.get_user(&principal).is_ok());
        }

        #[test]
        fn test_user_state_create_user(principal in principal_strategy(), user_args: UserDataArgs, address_args: UserAddressArgs) {
            let mut user_state = initialize_state();

            wallet_canister_caller();

            let user_data = user_state.create_user(principal, user_args.clone(), address_args.clone()).unwrap();

            assert_eq!(user_data.balance, user_args.balance.unwrap_or_default());
            assert_eq!(user_data.addresses[0].name, address_args.name.unwrap_or("Account 0".to_owned()));
        }

        #[test]
        fn test_user_state_get_user(principal in principal_strategy(), user_args: UserDataArgs, address_args: UserAddressArgs) {
            let mut user_state = initialize_state();

            wallet_canister_caller();

            user_state.create_user(principal, user_args.clone(), address_args.clone()).unwrap();

            assert!(user_state.get_user(&principal).is_ok());

            set_caller(principal);

            let user_data = user_state.get_user(&principal).unwrap();

            assert_eq!(user_data.balance, user_args.balance.unwrap_or_default());
            assert_eq!(user_data.addresses[0].name, address_args.name.unwrap_or("Account 0".to_owned()));
        }

        #[test]
        fn test_user_state_get_user_mut(principal in principal_strategy(), user_args: UserDataArgs, address_args: UserAddressArgs) {
            let mut user_state = initialize_state();

            wallet_canister_caller();

            user_state.create_user(principal, user_args.clone(), address_args.clone()).unwrap();

            set_caller(principal);

            assert!(user_state.get_user_mut(&principal).is_err());

            wallet_canister_caller();

            let user_data = user_state.get_user_mut(&principal).unwrap();

            assert_eq!(user_data.balance, user_args.balance.unwrap_or_default());
            assert_eq!(user_data.addresses[0].name, address_args.name.unwrap_or("Account 0".to_owned()));
        }
    }
}