#!/usr/bin/env python3

import argparse
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="Withdraw funds from a Mango account")
mango.ContextBuilder.add_command_line_parameters(parser)
mango.Wallet.add_command_line_parameters(parser)
parser.add_argument(
    "--symbol", type=str, required=True, help="token symbol to withdraw (e.g. USDC)"
)
parser.add_argument(
    "--quantity", type=Decimal, required=False, help="quantity token to withdraw"
)
parser.add_argument(
    "--wallet-target",
    type=Decimal,
    required=False,
    help="wallet balance of token to target with withdrawal",
)
parser.add_argument(
    "--account-target",
    type=Decimal,
    required=False,
    help="Mango account balance of token to target with withdrawal",
)
parser.add_argument(
    "--account-address",
    type=PublicKey,
    help="address of the specific account to use, if more than one available",
)
parser.add_argument(
    "--destination-wallet",
    type=PublicKey,
    help="if specified, the wallet to which the withdrawal should be sent. (Defaults to the current wallet.)",
)
parser.add_argument(
    "--allow-borrow",
    action="store_true",
    default=False,
    help="allow borrowing to fund the withdrawal",
)
parser.add_argument(
    "--wait",
    action="store_true",
    default=False,
    help="wait until the transactions are confirmed",
)
parser.add_argument(
    "--dry-run",
    action="store_true",
    default=False,
    help="runs as read-only and does not perform any transactions",
)
args: argparse.Namespace = mango.parse_args(parser)

if args.quantity is None and args.wallet_target is None and args.account_target is None:
    raise Exception(
        "One (and only one) of --quantity, --wallet-target or --account-target must be specified"
    )


def __quantity_from_account_target(
    account: mango.Account,
    token: mango.Token,
    quantity: typing.Optional[Decimal],
    allow_borrow: bool,
) -> typing.Optional[Decimal]:
    if quantity is None:
        return None

    slot = account.slot_by_instrument(token)
    to_withdraw = slot.net_value.value - quantity
    if to_withdraw < 0:
        if allow_borrow:
            to_withdraw = 0 - to_withdraw
        else:
            raise Exception(
                f"Cannot achieve Mango Account target of {quantity:,.8f} {token.symbol} by withdrawing - Mango Account only has balance of {slot.net_value.value:,.8f} {token.symbol}"
            )

    rounded = token.round(to_withdraw, mango.RoundDirection.DOWN)
    mango.output(
        f"Withdrawing {rounded:,.8f} {token.symbol} from Mango Account balance of {slot.net_value.value:,.8f} {token.symbol}"
    )
    return rounded


def __quantity_from_wallet_target(
    context: mango.Context,
    wallet: mango.Wallet,
    account: mango.Account,
    token: mango.Token,
    quantity: typing.Optional[Decimal],
    allow_borrow: bool,
) -> typing.Optional[Decimal]:
    if quantity is None:
        return None

    token_accounts: typing.Sequence[
        mango.TokenAccount
    ] = mango.TokenAccount.fetch_all_for_owner_and_token(context, wallet.address, token)
    total = sum(acc.value.value for acc in token_accounts)
    to_withdraw = quantity - total
    if to_withdraw < 0:
        raise Exception(
            f"Cannot achieve Mango Account target of {quantity:,.8f} {token.symbol} by withdrawing - wallet already has balance of {total:,.8f} {token.symbol}"
        )

    slot = account.slot_by_instrument(token)
    borrow = to_withdraw - slot.net_value.value
    if borrow > 0 and not allow_borrow:
        raise Exception(
            f"Cannot achieve Mango Account target of {quantity:,.8f} {token.symbol} without borrowing {borrow:,.8f} {token.symbol} (rerun specifying --allow-borrow if you want to borrow enough to achieve this wallet target)."
        )

    rounded = token.round(to_withdraw, mango.RoundDirection.DOWN)
    mango.output(
        f"Withdrawing {rounded:,.8f} {token.symbol} from Mango Account (borrowing {borrow:,.8f} {token.symbol}) to add to wallet balance of {total:,.8f} {token.symbol}"
    )
    return rounded


with mango.ContextBuilder.from_command_line_parameters(args) as context:
    wallet = mango.Wallet.from_command_line_parameters_or_raise(args)

    group = mango.Group.load(context, context.group_address)
    account = mango.Account.load_for_owner_by_address(
        context, wallet.address, group, args.account_address
    )

    token = mango.token(context, args.symbol)
    quantity: typing.Optional[Decimal] = (
        args.quantity
        or __quantity_from_account_target(
            account, token, args.account_target, args.allow_borrow
        )
        or __quantity_from_wallet_target(
            context, wallet, account, token, args.wallet_target, args.allow_borrow
        )
    )
    if quantity is None:
        raise Exception(
            "None of --quantity, --wallet-target or --account-target were specified - must specify one (and only one) of those parameters"
        )

    if quantity < 0:
        raise Exception(
            f"Cannot withdraw negative quantity {quantity:,.8f} {token.symbol}"
        )

    if quantity == 0:
        mango.output("Quantity to withdraw is 0 - nothing to do")
    else:
        withdrawal_value = mango.InstrumentValue(token, quantity)
        if args.dry_run:
            mango.output("Dry run - not sending transaction")
        else:
            destination = args.destination_wallet or wallet.address
            signatures = account.withdraw(
                context, wallet, destination, withdrawal_value, args.allow_borrow
            )

            if args.wait:
                mango.output("Waiting on transaction signatures:")
                mango.output(mango.indent_collection_as_str(signatures, 1))
                results = mango.WebSocketTransactionMonitor.wait_for_all(
                    context.client.cluster_ws_url, signatures
                )
                mango.output("Transaction results:")
                mango.output(mango.indent_collection_as_str(results, 1))
            else:
                mango.output("Transaction signatures:")
                mango.output(mango.indent_collection_as_str(signatures, 1))
