#!/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_redeem_accrued_mango_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 transaction is 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")

context = mango.ContextBuilder.from_command_line_parameters(args)
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:
    market = context.market_lookup.find_by_symbol(args.market)
    if market is None:
        raise Exception(f"Could not find market {args.market}")

    loaded_market: mango.LoadedMarket = mango.ensure_market_loaded(context, market)
    if not isinstance(loaded_market, mango.PerpMarket):
        raise Exception(f"Market {args.market} is not a perp market")
    perp_market = loaded_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)

transaction_ids = all_instructions.execute(context)
mango.output("Transaction IDs:", transaction_ids)

if args.wait:
    context.client.wait_for_confirmation(transaction_ids)
    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)
