#!/usr/bin/env pyston3

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

from decimal import Decimal

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="Balance the value of tokens in a Mango Markets group to specific values or percentages.")
mango.ContextBuilder.add_command_line_parameters(parser)
mango.Wallet.add_command_line_parameters(parser)
parser.add_argument("--target", type=mango.parse_fixed_target_balance, action="append", required=True,
                    help="token symbol plus target value, separated by a colon (e.g. 'ETH:2.5')")
parser.add_argument("--action-threshold", type=Decimal, default=Decimal("0.01"),
                    help="fraction of total wallet value a trade must be above to be carried out")
parser.add_argument("--adjustment-factor", type=Decimal, default=Decimal("0.05"),
                    help="factor by which to adjust the SELL price (akin to maximum slippage)")
parser.add_argument("--quote-symbol", type=str, default="USDC", help="quote token symbol to use for markets")
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.Context = mango.ContextBuilder.from_command_line_parameters(args)
wallet: mango.Wallet = mango.Wallet.from_command_line_parameters_or_raise(args)
action_threshold: Decimal = args.action_threshold
adjustment_factor: Decimal = args.adjustment_factor

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

targets: typing.Sequence[mango.FixedTargetBalance] = args.target
logging.info(f"Targets: {targets}")

if args.dry_run:
    trade_executor: mango.TradeExecutor = mango.NullTradeExecutor()
else:
    trade_executor = mango.ImmediateTradeExecutor(context, wallet, None, adjustment_factor)

quote_token: typing.Optional[mango.Token] = context.token_lookup.find_by_symbol(args.quote_symbol)
if quote_token is None:
    raise Exception(f"Could not find quote token '{args.quote_symbol}.")

prices: typing.List[mango.TokenValue] = []
oracle_provider: mango.OracleProvider = mango.create_oracle_provider(context, "market")
for target in targets:
    target_token: typing.Optional[mango.Token] = context.token_lookup.find_by_symbol(target.symbol)
    if target_token is None:
        raise Exception(f"Could not find target token '{target.symbol}.")
    market_symbol: str = f"serum:{target_token.symbol}/{quote_token.symbol}"
    market = context.market_lookup.find_by_symbol(market_symbol)
    if market is None:
        raise Exception(f"Could not find market {market_symbol}")
    oracle = oracle_provider.oracle_for_market(context, market)
    if oracle is None:
        raise Exception(f"Could not find oracle for market {market_symbol}")
    price = oracle.fetch_price(context)
    prices += [mango.TokenValue(target_token, price.mid_price)]
prices += [mango.TokenValue(quote_token, Decimal(1))]

wallet_balancer = mango.LiveWalletBalancer(wallet, quote_token, trade_executor, targets, action_threshold)
wallet_balancer.balance(context, prices)

logging.info("Balancing completed.")