class StakingPool:
    def __init__(self):
        # TODO: Add mutex
        # self.mutex = None

        self.fee_rate = 0.03

        self.share_pct = 0.6
        self.community_pct = 0.3
        self.developer_pct = 0.1

        self.contract_fund = 0.0
        self.community_fund = 0.0
        self.developer_fund = 0.0

        self.user_balances = {}

    def stake(self, user: str, amount: float):
        # Calculate fee
        total_fee = amount * self.fee_rate

        # Distribute fee
        # 1. Community funds
        community_fee = total_fee * self.community_pct
        self.community_fund += community_fee

        # 2. Developer funds
        developer_fee = total_fee * self.developer_pct
        self.developer_fund += developer_fee

        # 3. Reward stakers
        # If there are no users in the pool, the reward is placed in the community fund
        reward_amount = total_fee - community_fee - developer_fee

        ########
        # Test no community fund (distribute equally)
        self.community_distribute()

        if len(self.user_balances) < 1:
            self.community_fund += reward_amount
        else:
            self.__distribute_reward(reward_amount)

        # Deposit amount minus fee to contract_fund (wallet)
        self.contract_fund += amount - total_fee

        # Update user_balances with deposit amount
        if user not in self.user_balances:
            self.user_balances[user] = amount - total_fee
        else:
            self.user_balances[user] += amount - total_fee

    def unstake(self, user: str):
        # Fetch amount from user_balances
        if user not in self.user_balances:
            return 0.0
        elif self.user_balances[user] <= 0:
            return 0.0

        user_balance = self.user_balances[user]

        # Remove user from pool to prevent self-reward leaving a hanging balance
        del self.user_balances[user]

        # Calculate fee
        total_fee = user_balance * self.fee_rate
        # Adjust working balance
        user_balance -= total_fee

        # Distribute fee
        # 1. Community funds
        community_fee = total_fee * self.community_pct
        self.community_fund += community_fee

        #########
        # TEST DISTRIBUTE TO COMMUNITY
        self.community_distribute()

        # 2. Developer funds
        developer_fee = total_fee * self.developer_pct
        self.developer_fund += developer_fee

        # 3. Reward stakers
        reward_amount = total_fee - community_fee - developer_fee
        self.__distribute_reward(reward_amount)

        # Transfer tokens from contract_funds (wallet) to user's wallet address
        return user_balance

    def dev_withdrawal(self, destination: str, amount: float):
        pass

    def community_withdrawal(self, destination: str, amount: float):
        pass

    def community_distribute(self):
        # Take all funds in community stash and distribute it to stakers
        amount = self.community_fund
        self.__distribute_reward(amount)
        self.community_fund = 0.0

    def __distribute_reward(self, amount):
        # Get total staked amount
        staked_total = sum(self.user_balances.values())

        # Determine distribution amongst users
        # Update self.user_balances
        for user, user_stake in self.user_balances.items():
            pct = user_stake / staked_total
            user_reward = pct * amount
            self.user_balances[user] += user_reward

    def stats(self):
        print('\n--- POOL STATS ---')
        print(f'staked:         {sum(self.user_balances.values())}')
        print(f'community fund: {self.community_fund}')
        print(f'developer fund: {self.developer_fund}')
        print('------------------')
        print('--- USER TOTALS ---')
        for user, balance in self.user_balances.items():
            print(f'{user}: {balance}')
        print('------------------')


def main():
    pool = StakingPool()
    pool.stake("owner", 50000)
    pool.stake("user1", 50000000)
    # pool.stake("user2", 5000)
    # pool.stake("user3", 50000)
    # pool.stake("user3", 50000)
    # pool.stake("user3", 50000)
    # pool.stake("user3", 50000)
    # pool.stake("user3", 50000)
    # pool.stake("user3", 50000)
    pool.stats()
    # pool.community_distribute()

    print(pool.unstake("owner"))
    pool.stats()


if __name__ == '__main__':
    main()