#!/usr/bin/env python3

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

from decimal import Decimal
from solana.publickey import PublicKey
from solana.rpc.types import TxOpts
from spl.token.client import Token
from spl.token.constants import ACCOUNT_LEN, TOKEN_PROGRAM_ID

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

# We explicitly want argument parsing to be outside the main try-except block because some arguments
# (like --help) will cause an exit, which our except: block traps.
parser = argparse.ArgumentParser(description="Sends SPL tokens to a different address.")
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 send (e.g. ETH)")
parser.add_argument("--address", type=PublicKey,
                    help="Destination address for the SPL token - can be either the actual token address or the address of the owner of the token address")
parser.add_argument("--quantity", type=Decimal, required=True, help="quantity of token to send")
parser.add_argument("--wait", action="store_true", default=False,
                    help="wait until the transaction is 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)

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

logging.info(f"Wallet address: {wallet.address}")

token = context.token_lookup.find_by_symbol(args.symbol.upper())
if token is None:
    raise Exception(f"Could not find details of token with symbol {args.symbol}.")

spl_token = Token(context.client.compatible_client, token.mint, TOKEN_PROGRAM_ID, wallet.keypair)
source_accounts = spl_token.get_accounts(wallet.address)
source_account = source_accounts["result"]["value"][0]
source = PublicKey(source_account["pubkey"])

# Is the address an actual token account? Or is it the SOL address of the owner?
account_info: typing.Optional[mango.AccountInfo] = mango.AccountInfo.load(context, args.address)
if account_info is None:
    raise Exception(f"Could not find account at address {args.address}.")

destination: PublicKey
if account_info.owner == mango.SYSTEM_PROGRAM_ADDRESS:
    # This is a root wallet account - get the token account to use.
    destination = mango.TokenAccount.find_or_create_token_address_to_use(context, wallet, args.address, token)
elif account_info.owner == TOKEN_PROGRAM_ID and len(account_info.data) == ACCOUNT_LEN:
    # This is not a root wallet account, this is an SPL token account.
    destination = args.address
else:
    raise Exception(f"Account {args.address} is neither a root wallet account nor an SPL token account.")

owner = wallet.keypair
amount = int(args.quantity * Decimal(10 ** token.decimals))

print("Balance:", source_account["account"]["data"]["parsed"]
      ["info"]["tokenAmount"]["uiAmountString"], token.name)
text_amount = f"{amount} {token.name} (@ {token.decimals} decimal places)"
print(f"Sending {text_amount}")
print(f"    From: {source}")
print(f"      To: {destination}")

if args.dry_run:
    print("Skipping actual transfer - dry run.")
else:
    transfer_response = spl_token.transfer(source, destination, owner, amount,
                                           opts=TxOpts(preflight_commitment=context.client.commitment))
    transaction_ids = [transfer_response["result"]]
    print(f"Transaction IDs: {transaction_ids}")
    if args.wait:
        context.client.wait_for_confirmation(transaction_ids)

        updated_balance = spl_token.get_balance(source)
        updated_balance_text = updated_balance["result"]["value"]["uiAmountString"]
        print(f"{text_amount} sent. Balance now: {updated_balance_text} {token.name}")
    else:
        print(f"{text_amount} sent.")