#!/usr/bin/env python3

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

import spl.token.instructions as spl_token

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

parser = argparse.ArgumentParser(description="mint SPL tokens to your wallet")
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 ensure the associated token account exists (e.g. USDC)",
)
parser.add_argument(
    "--wait",
    action="store_true",
    default=False,
    help="wait until the transactions are confirmed",
)
args: argparse.Namespace = mango.parse_args(parser)

with mango.ContextBuilder.from_command_line_parameters(args) as context:
    wallet = mango.Wallet.from_command_line_parameters_or_raise(args)
    token: mango.Token = mango.token(context, args.symbol)

    associated_token_address = spl_token.get_associated_token_address(
        wallet.address, token.mint
    )
    token_account: typing.Optional[mango.TokenAccount] = mango.TokenAccount.load(
        context, associated_token_address
    )
    if token_account is not None:
        # The associated token account exists
        mango.output(
            f"Associated token account already exists at: {associated_token_address}."
        )
    else:
        # Create the proper associated token account.
        signer = mango.CombinableInstructions.from_wallet(wallet)
        create_instruction = spl_token.create_associated_token_account(
            wallet.address, wallet.address, token.mint
        )
        create = mango.CombinableInstructions.from_instruction(create_instruction)

        mango.output(
            f"No associated token account at: {associated_token_address} - creating..."
        )
        signatures = (signer + create).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))
        else:
            mango.output("Transaction signatures:")
            mango.output(mango.indent_collection_as_str(signatures, 1))

        mango.output(
            f"Associated token account created at: {associated_token_address}."
        )
