#!/usr/bin/env python3

import argparse
import logging
import os
import os.path
import sys
import typing

from decimal import Decimal
from solana.publickey import PublicKey

sys.path.insert(0, os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..")))
import mango  # nopep8

parser = argparse.ArgumentParser(description="Display the balances of all group tokens in the current wallet.")
mango.ContextBuilder.add_command_line_parameters(parser)
mango.Wallet.add_command_line_parameters(parser)
parser.add_argument("--address", type=PublicKey,
                    help="Root address to check (if not provided, the wallet address is used)")
args: argparse.Namespace = mango.parse_args(parser)

address: typing.Optional[PublicKey] = args.address
if address is None:
    wallet = mango.Wallet.from_command_line_parameters_or_raise(args)
    address = wallet.address

context: mango.Context = mango.ContextBuilder.from_command_line_parameters(args)
group: mango.Group = mango.Group.load(context)
cache: mango.Cache = mango.Cache.load(context, group.cache)

logging.info(f"Address: {address}")

balances = group.fetch_balances(context, address)
print(f"\nToken Balances [{address}]:")
total_in_wallet: mango.TokenValue = mango.TokenValue(group.shared_quote_token.token, Decimal(0))
for balance in balances:
    if balance.value != 0:
        balance_text: str = f"{balance} "
        value_text: str = ""
        if balance.token == group.shared_quote_token.token:
            total_in_wallet = total_in_wallet + balance
            value_text = f" worth {balance}"
        else:
            token_index: typing.Optional[int] = group.find_token_market_index_or_none(balance.token)
            if token_index is not None:
                cached_token_price: typing.Optional[mango.PriceCache] = cache.price_cache[token_index]
                if cached_token_price is not None:
                    balance_value: mango.TokenValue = mango.TokenValue(
                        group.shared_quote_token.token, balance.value * cached_token_price.price)
                    total_in_wallet += balance_value
                    value_text = f" worth {balance_value}"
        print(f"               {balance_text:<45}{value_text}")
print(
    f"Total Value:   {total_in_wallet}")


mango_accounts = mango.Account.load_all_for_owner(context, address, group)
account_value: mango.TokenValue = mango.TokenValue(group.shared_quote_token.token, Decimal(0))
quote_token_free_in_open_orders: mango.TokenValue = mango.TokenValue(group.shared_quote_token.token, Decimal(0))
quote_token_total_in_open_orders: mango.TokenValue = mango.TokenValue(group.shared_quote_token.token, Decimal(0))
grand_total: mango.TokenValue = total_in_wallet
for account in mango_accounts:
    print(f"\nAccount Balances [{account.address}]:")
    at_least_one_output: bool = False
    open_orders: typing.Dict[str, mango.OpenOrders] = account.load_all_spot_open_orders(context)
    for asset in account.basket:
        if (asset.deposit.value != 0) or (asset.borrow.value != 0) or (asset.net_value.value != 0) or ((asset.perp_account is not None) and not asset.perp_account.empty):
            at_least_one_output = True
            report: mango.AccountTokenValues = mango.AccountTokenValues.from_account_basket_base_token(
                asset, open_orders, group)
            # print(report)
            price_from_cache: mango.TokenValue = group.token_price_from_cache(cache, report.base_token)
            perp_market_cache: typing.Optional[mango.PerpMarketCache] = group.perp_market_cache_from_cache(
                cache, report.base_token)
            priced_report: mango.AccountTokenValues = report.priced(price_from_cache, perp_market_cache)
            account_value += priced_report.net_value
            quote_token_free_in_open_orders += priced_report.quote_token_free
            quote_token_total_in_open_orders += priced_report.quote_token_total
            print(priced_report)

    quote_report: mango.AccountTokenValues = mango.AccountTokenValues(account.shared_quote_token.token_info.token,
                                                                      account.shared_quote_token.token_info.token,
                                                                      account.shared_quote_token.deposit,
                                                                      account.shared_quote_token.borrow,
                                                                      mango.TokenValue(
                                                                          group.shared_quote_token.token, Decimal(0)),
                                                                      mango.TokenValue(
                                                                          group.shared_quote_token.token, Decimal(0)),
                                                                      quote_token_free_in_open_orders,
                                                                      quote_token_total_in_open_orders,
                                                                      mango.TokenValue(
                                                                          group.shared_quote_token.token, Decimal(0)),
                                                                      Decimal(0), Decimal(0), Decimal(0),
                                                                      mango.NullLotSizeConverter())
    account_value += quote_report.net_value + quote_token_total_in_open_orders
    print(quote_report)

    print(f"Account Total: {account_value}")
    grand_total += account_value

print(f"\nGrand Total:   {grand_total}")