import { BatchCall } from '@aztec/aztec.js';
import chunk from 'lodash.chunk';
export class TokenSimulator {
    token;
    defaultWallet;
    defaultAddress;
    logger;
    accounts;
    balancesPrivate;
    balancePublic;
    totalSupply;
    lookupProvider;
    constructor(token, defaultWallet, defaultAddress, logger, accounts){
        this.token = token;
        this.defaultWallet = defaultWallet;
        this.defaultAddress = defaultAddress;
        this.logger = logger;
        this.accounts = accounts;
        this.balancesPrivate = new Map();
        this.balancePublic = new Map();
        this.totalSupply = 0n;
        this.lookupProvider = new Map();
    }
    addAccount(account) {
        this.accounts.push(account);
    }
    setLookupProvider(account, wallet) {
        this.lookupProvider.set(account.toString(), wallet);
    }
    mintPrivate(to, amount) {
        this.totalSupply += amount;
        this.balancesPrivate.set(to.toString(), (this.balancesPrivate.get(to.toString()) || 0n) + amount);
    }
    mintPublic(to, amount) {
        this.totalSupply += amount;
        const value = this.balancePublic.get(to.toString()) || 0n;
        this.balancePublic.set(to.toString(), value + amount);
    }
    transferPublic(from, to, amount) {
        const fromBalance = this.balancePublic.get(from.toString()) || 0n;
        this.balancePublic.set(from.toString(), fromBalance - amount);
        expect(fromBalance).toBeGreaterThanOrEqual(amount);
        const toBalance = this.balancePublic.get(to.toString()) || 0n;
        this.balancePublic.set(to.toString(), toBalance + amount);
    }
    transferPrivate(from, to, amount) {
        const fromBalance = this.balancesPrivate.get(from.toString()) || 0n;
        expect(fromBalance).toBeGreaterThanOrEqual(amount);
        this.balancesPrivate.set(from.toString(), fromBalance - amount);
        const toBalance = this.balancesPrivate.get(to.toString()) || 0n;
        this.balancesPrivate.set(to.toString(), toBalance + amount);
    }
    transferToPrivate(from, to, amount) {
        const fromBalance = this.balancePublic.get(from.toString()) || 0n;
        expect(fromBalance).toBeGreaterThanOrEqual(amount);
        this.balancePublic.set(from.toString(), fromBalance - amount);
        const toBalance = this.balancesPrivate.get(to.toString()) || 0n;
        this.balancesPrivate.set(to.toString(), toBalance + amount);
    }
    transferToPublic(from, to, amount) {
        const fromBalance = this.balancesPrivate.get(from.toString()) || 0n;
        const toBalance = this.balancePublic.get(to.toString()) || 0n;
        expect(fromBalance).toBeGreaterThanOrEqual(amount);
        this.balancesPrivate.set(from.toString(), fromBalance - amount);
        this.balancePublic.set(to.toString(), toBalance + amount);
    }
    burnPrivate(from, amount) {
        const fromBalance = this.balancesPrivate.get(from.toString()) || 0n;
        expect(fromBalance).toBeGreaterThanOrEqual(amount);
        this.balancesPrivate.set(from.toString(), fromBalance - amount);
        this.totalSupply -= amount;
    }
    burnPublic(from, amount) {
        const fromBalance = this.balancePublic.get(from.toString()) || 0n;
        expect(fromBalance).toBeGreaterThanOrEqual(amount);
        this.balancePublic.set(from.toString(), fromBalance - amount);
        this.totalSupply -= amount;
    }
    balanceOfPublic(address) {
        return this.balancePublic.get(address.toString()) || 0n;
    }
    balanceOfPrivate(address) {
        return this.balancesPrivate.get(address.toString()) || 0n;
    }
    async checkPublic() {
        // public calls
        const calls = [
            this.token.methods.total_supply(),
            ...this.accounts.map((address)=>this.token.methods.balance_of_public(address))
        ];
        const results = (await Promise.all(chunk(calls, 4).map((batch)=>new BatchCall(this.defaultWallet, batch).simulate({
                from: this.defaultAddress
            })))).flat();
        expect(results[0]).toEqual(this.totalSupply);
        // Check that all our balances match
        for(let i = 0; i < this.accounts.length; i++){
            expect(results[i + 1]).toEqual(this.balanceOfPublic(this.accounts[i]));
        }
    }
    async checkPrivate() {
        // Private calls
        const defaultLookups = [];
        const nonDefaultLookups = [];
        for (const address of this.accounts){
            if (this.lookupProvider.has(address.toString())) {
                nonDefaultLookups.push(address);
            } else {
                defaultLookups.push(address);
            }
        }
        const defaultCalls = defaultLookups.map((address)=>this.token.methods.balance_of_private(address));
        const results = (await Promise.all(chunk(defaultCalls, 4).map((batch)=>new BatchCall(this.defaultWallet, batch).simulate({
                from: this.defaultAddress
            })))).flat();
        for(let i = 0; i < defaultLookups.length; i++){
            expect(results[i]).toEqual(this.balanceOfPrivate(defaultLookups[i]));
        }
        // We are just running individual calls for the non-default lookups
        // @todo We should also batch these
        for (const address of nonDefaultLookups){
            const wallet = this.lookupProvider.get(address.toString());
            const asset = wallet ? this.token.withWallet(wallet) : this.token;
            const actualPrivateBalance = await asset.methods.balance_of_private({
                address
            }).simulate({
                from: this.defaultAddress
            });
            expect(actualPrivateBalance).toEqual(this.balanceOfPrivate(address));
        }
    }
    async check() {
        await this.checkPublic();
        await this.checkPrivate();
    }
}
