#!/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


def report_accrued(basket_token: mango.AccountSlot) -> None:
    symbol: str = basket_token.base_instrument.symbol
    if basket_token.perp_account is None:
        accrued: mango.InstrumentValue = mango.InstrumentValue(
            basket_token.base_instrument, Decimal(0)
        )
    else:
        accrued = basket_token.perp_account.mngo_accrued
    mango.output(f"Accrued in perp market [{symbol:>5}]: {accrued}")


def load_perp_market(
    context: mango.Context, group: mango.Group, slot: mango.GroupSlot
) -> typing.Optional[mango.PerpMarket]:
    if slot.perp_market is None:
        return None

    perp_market_details = mango.PerpMarketDetails.load(
        context, slot.perp_market.address, group
    )
    perp_market = mango.PerpMarket(
        context.mango_program_address,
        slot.perp_market.address,
        slot.base_instrument,
        mango.Token.ensure(slot.quote_token_bank.token),
        perp_market_details,
    )

    return perp_market


def find_basket_token_in_account(
    account: mango.Account, instrument: mango.Instrument
) -> typing.Optional[mango.AccountSlot]:
    basket_tokens = [
        in_basket
        for in_basket in account.slots
        if in_basket.base_instrument == instrument
    ]

    if len(basket_tokens) == 0:
        return None
    else:
        return basket_tokens[0]


def build_redeem_instruction_for_account(
    context: mango.Context,
    wallet: mango.Wallet,
    group: mango.Group,
    mngo: mango.TokenBank,
    account: mango.Account,
    perp_market: mango.PerpMarket,
    basket_token: typing.Optional[mango.AccountSlot],
) -> mango.CombinableInstructions:
    if (
        (basket_token is None)
        or (basket_token.perp_account is None)
        or basket_token.perp_account.mngo_accrued.value == 0
    ):
        return mango.CombinableInstructions.empty()

    report_accrued(basket_token)
    redeem = mango.build_mango_redeem_accrued_instructions(
        context, wallet, perp_market, group, account, mngo
    )

    return redeem


parser = argparse.ArgumentParser(
    description="redeems accrued MNGO from a Mango account"
)
mango.ContextBuilder.add_command_line_parameters(parser)
mango.Wallet.add_command_line_parameters(parser)
parser.add_argument(
    "--market", type=str, help="perp market symbol with accrued MNGO (e.g. ETH-PERP)"
)
parser.add_argument(
    "--all",
    action="store_true",
    default=False,
    help="redeem all MNGO in all perp markets in the account",
)
parser.add_argument(
    "--account-address",
    type=PublicKey,
    help="address of the specific account to use, if more than one available",
)
parser.add_argument(
    "--wait",
    action="store_true",
    default=False,
    help="wait until the transactions are confirmed",
)
args: argparse.Namespace = mango.parse_args(parser)

if (not args.all) and (args.market is None):
    raise Exception(
        "Must specify either an individual market (using --market) or use --all for all markets"
    )

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)
    mngo = group.liquidity_incentive_token_bank
    account = mango.Account.load_for_owner_by_address(
        context, wallet.address, group, args.account_address
    )

    signers: mango.CombinableInstructions = mango.CombinableInstructions.from_wallet(
        wallet
    )
    all_instructions: mango.CombinableInstructions = signers

    if args.all:
        for slot in group.slots:
            perp_market = load_perp_market(context, group, slot)
            if perp_market is not None:
                basket_token = find_basket_token_in_account(
                    account, slot.base_instrument
                )
                all_instructions += build_redeem_instruction_for_account(
                    context, wallet, group, mngo, account, perp_market, basket_token
                )
    else:
        perp_market = mango.PerpMarket.ensure(mango.market(context, args.market))
        basket_token = find_basket_token_in_account(account, perp_market.base)
        all_instructions += build_redeem_instruction_for_account(
            context, wallet, group, mngo, account, perp_market, basket_token
        )

    signatures = all_instructions.execute(context)

    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))

        reloaded_account = mango.Account.load_for_owner_by_address(
            context, wallet.address, group, args.account_address
        )
        if args.all:
            for slot in group.slots:
                basket_token = find_basket_token_in_account(
                    reloaded_account, slot.base_instrument
                )
                if basket_token is not None:
                    report_accrued(basket_token)
        elif perp_market is not None:
            basket_token = find_basket_token_in_account(
                reloaded_account, perp_market.base
            )
            if basket_token is not None:
                report_accrued(basket_token)
    else:
        mango.output("Transaction signatures:")
        mango.output(mango.indent_collection_as_str(signatures, 1))
