#!/usr/bin/env python3

import argparse
import os
import os.path
import pandas
import sys
import typing

from solana.publickey import PublicKey

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

parser = argparse.ArgumentParser(
    description="Display the balances of all group tokens in the current wallet."
)
entropy.ContextBuilder.add_command_line_parameters(parser)
entropy.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 = entropy.parse_args(parser)

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

with entropy.ContextBuilder.from_command_line_parameters(args) as context:
    group: entropy.Group = entropy.Group.load(context)
    cache: entropy.Cache = entropy.Cache.load(context, group.cache)

    address_account_info: typing.Optional[
        entropy.AccountInfo
    ] = entropy.AccountInfo.load(context, address)
    if address_account_info is None:
        raise Exception(f"Could not load account data from address {address}")

    entropy_accounts: typing.Sequence[entropy.Account]
    if len(address_account_info.data) == entropy.layouts.MANGO_ACCOUNT.sizeof():
        entropy_accounts = [entropy.Account.parse(address_account_info, group, cache)]
    else:
        entropy_accounts = entropy.Account.load_all_for_owner(context, address, group)

for account in entropy_accounts:
    pandas.set_option("display.max_columns", None)
    pandas.set_option("display.width", None)
    pandas.set_option("display.precision", 6)

    open_orders: typing.Dict[
        str, entropy.OpenOrders
    ] = account.load_all_spot_open_orders(context)
    frame: pandas.DataFrame = account.to_dataframe(group, open_orders, cache)
    entropy.output(frame)
    if entropy.output_formatter.format == entropy.OutputFormat.TEXT:
        entropy.output("Init Health:", account.init_health(frame))
        entropy.output("Maint Health:", account.maint_health(frame))
        entropy.output("Total Value:", account.total_value(frame))
