Move Complex Global Data Structuring

module jyunth_addr::MetalVault {
    use std::string::{String, utf8};
    use std::signer;
    use std::simple_map::{Self, SimpleMap};
    use std::debug::print;
  
    const GOLD: u64 = 0;
    const SILVER: u64 = 0;
  
    struct MetalReserves has store, copy, drop {
        g28: simple_map::SimpleMap<String, u64>,
        g57: simple_map::SimpleMap<String, u64>,
        g114: simple_map::SimpleMap<String, u64>
    }
  
    struct MetalVault has key {
        gold: MetalReserves,
        silver: MetalReserves,
    }
  
    fun init_client(account: &signer) {
        let metal_vault: SimpleMap<String, u64> = simple_map::create();
        simple_map::add(&mut metal_vault, utf8(b"UAE"), 0);
        simple_map::add(&mut metal_vault, utf8(b"MEX"), 0);
        simple_map::add(&mut metal_vault, utf8(b"COL"), 0);
  
        let init_balance = MetalReserves {
            g28: metal_vault,
            g57: metal_vault,
            g114: metal_vault
        };
  
        let vault = MetalVault {
            gold: init_balance,
            silver: init_balance
        };
  
        move_to(account, vault);
    }
  
    fun get_vault(account: address): (MetalReserves, MetalReserves) acquires MetalVault {
        (borrow_global<MetalVault>(account).gold, borrow_global<MetalVault>(account).silver)
    }
  
    fun read_balance(asset_bal: SimpleMap<String, u64>, grams: String) {
        print(&grams);
        print(&utf8(b"UAE"));
        print(simple_map::borrow(&mut asset_bal, &utf8(b"UAE")));
        print(&utf8(b"MEX"));
        print(simple_map::borrow(&mut asset_bal, &utf8(b"MEX")));
        print(&utf8(b"COL"));
        print(simple_map::borrow(&mut asset_bal, &utf8(b"COL")));
    }
  
    fun get_client_balance(account: address, asset: u64) acquires MetalVault {
        let (gold, silver) = get_vault(account);
        if (asset == GOLD) {
            read_balance(gold.g28, utf8(b"28 grams"));
            read_balance(gold.g57, utf8(b"57 grams"));
            read_balance(gold.g114, utf8(b"114 grams"));
        } else {
            read_balance(silver.g28, utf8(b"28 grams"));
            read_balance(silver.g57, utf8(b"57 grams"));
            read_balance(silver.g114, utf8(b"114 grams"));
        }
    }
  
    fun update_balance(metal: &mut MetalReserves, country: String, amount: u64, weight: u64): bool {
        if (weight == 28) {
            let current = simple_map::borrow_mut(&mut metal.g28, &country);
            *current = *current + amount;
            return true
        } else if (weight == 57) {
            let current = simple_map::borrow_mut(&mut metal.g57, &country);
            *current = *current + amount;
            return true
        } else if (weight == 114) {
            let current = simple_map::borrow_mut(&mut metal.g114, &country);
            *current = *current + amount;
            return true
        } else {
            return false
        }
    }
  
    fun add_metal(account: address, country: String, type: u64, amount: u64, weight: u64): bool acquires MetalVault {
        if (type == GOLD) {
            let metal = &mut borrow_global_mut<MetalVault>(account).gold;
            update_balance(metal, country, amount, weight);
            return true
        } else if (type == SILVER) {
            let metal = &mut borrow_global_mut<MetalVault>(account).silver;
            update_balance(metal, country, amount, weight);
            return true
        } else {
            return false
        }
    }
  
    #[test(client1 = @0x123)]
    fun test_function(client1: signer) acquires MetalVault {
        init_client(&client1);
        assert!(exists<MetalVault>(signer::address_of(&client1)) == true, 404);
        add_metal(signer::address_of(&client1), utf8(b"UAE"), GOLD, 3, 57);
        get_client_balance(signer::address_of(&client1), GOLD);
    }
}