from decimal import Decimal, getcontext

# Set precision context if needed
getcontext().prec = 28

class StakingContract:
    def __init__(self):
        self.rewards_pct = Decimal('0.7')
        self.community_pct = Decimal('0.2')
        self.developer_pct = Decimal('0.1')

        self.balances = {
            "stakes": Decimal('0.0'),
            "reward": Decimal('0.0'),
            "developer": Decimal('0.0'),
            "community": Decimal('0.0'),
        }

        self.users = {}

    def deposit(self, user: str, amount: float):
        amount = Decimal(amount)  # Convert to Decimal
        stake, reward, community, developer = self.calc_fee(amount)
        print(f'{user} staked {amount} ($JENNER)')

        self.add_fees_to_balances(stake, reward, community, developer)
        self.add_stake_to_user(user, stake)

    def withdraw(self, user: str):
        if user not in self.users:
            print(f"No staked amount found for user {user}.")
            return Decimal('0.0')

        initial_stake = self.users[user]
        reward_share = self.calculate_reward_share(user)

        # Update balances
        self.balances['stakes'] -= initial_stake
        self.balances['reward'] -= reward_share

        # Reset user stake to 0
        del self.users[user]

        # Total amount returned to the user
        total_withdrawal = initial_stake + reward_share
        print(f'{user} withdrew {total_withdrawal:.4f} ($JENNER) - (reward: {reward_share:.4f})')
        return total_withdrawal

    # def withdraw(self, user: str):
    #     if user not in self.users:
    #         print(f"No staked amount found for user {user}.")
    #         return Decimal('0.0')
    #
    #     initial_stake = self.users[user]
    #     reward_share = self.calculate_reward_share(user)
    #
    #     # Update balances
    #     self.balances['stakes'] -= initial_stake
    #     self.balances['reward'] -= reward_share
    #
    #     # Reset user stake to 0
    #     # self.users[user] = Decimal('0.0')
    #     del self.users[user]
    #
    #     # Total amount returned to the user
    #     total_withdrawal = initial_stake + reward_share
    #     print(f'{user} withdrew {total_withdrawal:.4f} ($JENNER) - (reward: {reward_share:.4f})')
    #     return total_withdrawal

    def calculate_reward_share(self, user):
        total_stakes = sum(self.users.values())
        if total_stakes == 0:  # Avoid division by zero
            return Decimal('0.0')
        user_stake = self.users[user]
        reward_share = (user_stake / total_stakes) * self.balances['reward']
        return reward_share

    def calc_fee(self, amount):
        fee = amount * Decimal('0.03')
        stake_amount = amount - fee

        rewards = fee * self.rewards_pct
        community = fee * self.community_pct
        developer = fee * self.developer_pct

        return stake_amount, rewards, community, developer

    def show_stats(self):
        print('\n--- POOL STATS ---')
        for k, v in self.balances.items():
            print(f'{k}: {v}')
        print('')
        print('--- USERS IN POOL ---')
        for k, v in self.users.items():
            print(f'{k}: {v}')
        print('')

    def add_stake_to_user(self, user: str, stake: Decimal):
        if user in self.users:
            self.users[user] += stake
        else:
            self.users[user] = stake

    def add_fees_to_balances(self, stake: Decimal, reward: Decimal, community: Decimal, developer: Decimal):
        self.balances['stakes'] += stake
        self.balances['reward'] += reward
        self.balances['developer'] += developer
        self.balances['community'] += community


def main():
    pool = StakingContract()

    pool.deposit("owner", 50000)
    pool.deposit("user1", 50000)
    pool.deposit("user2", 50000)

    pool.show_stats()

    pool.withdraw("user1")
    pool.withdraw("user2")
    pool.withdraw("owner")

    pool.show_stats()

    # pool = StakingContract()
    #
    # pool.deposit("owner", 50000)
    # pool.deposit("user5", 10000000)
    # pool.deposit("user1", 1000)
    # pool.deposit("user2", 200000)
    # pool.deposit("user3", 90283)
    # pool.deposit("user4", 1000)
    # pool.deposit("user5", 1000000)
    # pool.deposit("user6", 1000000)
    # # pool.deposit("user5", 1000000)
    # # pool.deposit("user5", 1000000)
    # # pool.deposit("user5", 1000000)
    # #
    # # for n in range(200):
    # #     pool.deposit(f"user{n+6}", 5000)
    # #
    # # pool.deposit("user5", 1000000)
    #
    # pool.show_stats()
    #
    # pool.withdraw("user1")
    # pool.withdraw("user2")
    # pool.withdraw("user3")
    # pool.withdraw("user4")
    # pool.withdraw("user5")
    # pool.withdraw("user6")
    # #
    # # for n in range(200):
    # #     pool.withdraw(f"user{n+6}")
    #
    # pool.withdraw("owner")
    #
    # pool.show_stats()


if __name__ == '__main__':
    main()