dftool

#!/usr/bin/env python

import os
import sys

import brownie
from enforce_typing import enforce_types

from util import (
    blockrange,
    blocktime,
    calcrewards,
    networkutil,
    csvs,
    dispense,
    getrate,
    query,
)
from util.base18 import fromBase18
from util.constants import BROWNIE_PROJECT as B
from util.networkutil import DEV_CHAINID
from util.oceanutil import recordDeployedContracts, OCEANtoken

CHAINID_EXAMPLES = (
    f"{DEV_CHAINID} for development, 1 for (eth) mainnet, 137 for polygon"
)

# ========================================================================
HELP_MAIN = f"""Data Farming tool, for use by OPF.

OPF steps:
1. query - run per CHAINID ({DEV_CHAINID}, 1, 137, ..)
2. getrate - run per basetoken (OCEAN, H2O)
3. calc - run per rewardtoken (store amt per CHAINID): OCEAN (1+137), EWT (246)..
4. dispense - run per rewardtoken*CHAINID: OCEAN*1, OCEAN*137, EWT*246..

Then LP steps:
5. call DFRewards.claim(), or related

Usage: dftool query|getrate|calc|dispense|..

  dftool query ST FIN NSAMP CSV_DIR CHAINID - query chain, get stakes & vols
  dftool getrate TOKEN_SYMBOL ST FIN CSV_DIR - get exchange rate
  dftool calc CSV_DIR TOT_TOKEN TOKEN_SYMBOL - from stakes/vols/rates, calc rewards
  dftool dispense CSV_DIR CHAINID [DFREWARDS_ADDR] [TOKEN_ADDR] - from rewards, dispense funds

  dftool querymany ST FIN NSAMP CSV_DIR CHAINID1 [CHAINID2 [CHAINID3 ..]] 
  dftool compile - compile contracts
  dftool manyrandom - deploy many random tokens & pools then consume (for testing)
  dftool newdfrewards CHAINID - deploy new DFRewards contract
  dftool mine BLOCKS [TIMEDELTA] - force chain to pass time (ganache only)
  dftool newacct - generate new account
  dftool newtoken CHAINID - generate new token (for testing)

  dftool acctinfo CHAINID ACCOUNT_ADDR [TOKEN_ADDR] - info about an account
  dftool chaininfo CHAINID - info about a network
  dftool help - this message

Transactions are signed with envvar 'DFTOOL_KEY`.
"""


@enforce_types
def do_help():
    print(HELP_MAIN)
    sys.exit(0)


# ========================================================================
@enforce_types
def do_query():
    HELP = f"""Query chain, output stakes & vols csvs

Usage: dftool query ST FIN NSAMP CSV_DIR CHAINID
  ST -- first block # to calc on | YYYY-MM-DD | YYYY-MM-DD_HH:MM
  FIN -- last block # to calc on | YYYY-MM-DD | YYYY-MM-DD_HH:MM | latest
  NSAMP -- # blocks to sample liquidity from, from blocks [ST, ST+1, .., FIN]
  CSV_DIR -- output dir for stakes-CHAINID.csv and poolvols-CHAINID.csv
  CHAINID -- {CHAINID_EXAMPLES}

Uses these envvars:
ADDRESS_FILE -- eg for barge: export ADDRESS_FILE={networkutil.chainIdToAddressFile(chainID=DEV_CHAINID)}
"""
    if len(sys.argv) not in [2 + 5]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "query"
    ST, FIN, NSAMP = sys.argv[2], sys.argv[3], int(sys.argv[4])
    CSV_DIR = sys.argv[5]
    CHAINID = int(sys.argv[6])
    print("dftool query: Begin")
    print(
        f"Arguments:\n "
        f"\n ST={ST}\n FIN={FIN}\n NSAMP={NSAMP}"
        f"\n CSV_DIR={CSV_DIR}"
        f"\n CHAINID={CHAINID}"
    )

    # extract envvars
    ADDRESS_FILE = _getAddressEnvvarOrExit()

    # shouldn't already have the output files
    c1 = csvs.stakesCsvFilename(CSV_DIR, CHAINID)
    if os.path.exists(c1):
        print(f"\nFile {c1} exists. Exiting.")
        sys.exit(0)
    c2 = csvs.poolvolsCsvFilename(CSV_DIR, CHAINID)
    if os.path.exists(c2):
        print(f"\nFile {c2} exists. Exiting.")
        sys.exit(0)

    # make directory if needed
    if not os.path.exists(CSV_DIR):
        print(f"Directory {CSV_DIR} did not exist, so created it")
        os.mkdir(CSV_DIR)

    # brownie setup
    networkutil.connect(CHAINID)
    chain = brownie.network.chain

    if "-" in ST:
        st_block = blocktime.timestrToBlock(chain, ST)
    else:
        st_block = int(ST)

    if FIN == "latest":
        fin_block = len(chain)
    elif "-" in FIN:
        fin_block = blocktime.timestrToBlock(chain, FIN)
    else:
        fin_block = int(FIN)

    # main work
    recordDeployedContracts(ADDRESS_FILE)
    seed = fin_block
    rng = blockrange.BlockRange(st_block, fin_block, NSAMP, seed)
    (Pi, Si, Vi) = query.query_all(rng, CHAINID)
    csvs.savePoolinfoCsv(Pi, Si, Vi, CSV_DIR, CHAINID)
    csvs.saveStakesCsv(Si, CSV_DIR, CHAINID)
    csvs.savePoolvolsCsv(Vi, CSV_DIR, CHAINID)

    print("dftool query: Done")


# ========================================================================
@enforce_types
def do_getrate():
    HELP = """Get exchange rate, and output rate csv.

Usage: dftool getrate TOKEN_SYMBOL ST FIN CSV_DIR
  TOKEN_SYMBOL -- e.g. OCEAN, H2O
  ST -- start time -- YYYY-MM-DD
  FIN -- end time -- YYYY-MM-DD
  CSV_DIR -- output directory for rate-TOKEN_SYMBOL.csv file
"""
    if len(sys.argv) not in [2 + 4]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "getrate"
    TOKEN_SYMBOL = sys.argv[2]
    ST, FIN = sys.argv[3], sys.argv[4]
    CSV_DIR = sys.argv[5]
    print(f"Arguments: ST={ST}, FIN={FIN}, CSV_DIR={CSV_DIR}")

    # shouldn't already have the output file
    c1 = csvs.rateCsvFilename(TOKEN_SYMBOL, CSV_DIR)
    if os.path.exists(c1):
        print(f"\nFile {c1} exists. Exiting.")
        sys.exit(0)

    # make directory if needed
    if not os.path.exists(CSV_DIR):
        print(f"Directory {CSV_DIR} did not exist, so created it")
        os.mkdir(CSV_DIR)

    # main work
    rate = getrate.getrate(TOKEN_SYMBOL, ST, FIN)
    print(f"rate = ${rate:.4f} / {TOKEN_SYMBOL}")
    csvs.saveRateCsv(TOKEN_SYMBOL, rate, CSV_DIR)


# ========================================================================
@enforce_types
def do_calc():
    HELP = """From stakes/vols/rates csvs, output rewards csvs.

Usage: dftool calc CSV_DIR TOT_TOKEN TOKEN_SYMBOL
  CSV_DIR -- directory: input stakes & volumes csvs, output rewards.csv
  TOT_TOKEN -- total amount of TOKEN to distribute (decimal, not wei)
  TOKEN_SYMBOL -- e.g. OCEAN
"""

    if len(sys.argv) not in [2 + 3]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "calc"
    CSV_DIR = sys.argv[2]
    TOT_TOKEN = float(sys.argv[3])
    TOKEN_SYMBOL = sys.argv[4]
    print(
        f"Arguments:\n TOT_TOKEN={TOT_TOKEN}, CSV_DIR={CSV_DIR}"
        f", TOKEN_SYMBOL={TOKEN_SYMBOL}"
    )

    # do we have the input files?
    if not csvs.stakesCsvFilenames(CSV_DIR):
        print(f"\nNo 'stakes*.csv' files in '{CSV_DIR}'. Exiting.")
        sys.exit(0)
    if not csvs.poolvolsCsvFilenames(CSV_DIR):
        print(f"\nNo 'pool_vol*.csv' files in '{CSV_DIR}'. Exiting.")
        sys.exit(0)
    if not csvs.rateCsvFilenames(CSV_DIR):
        print(f"\nNo 'rate*.csv' files in '{CSV_DIR}'. Exiting.")
        sys.exit(0)

    # shouldn't already have the output file
    c1 = csvs.rewardsperlpCsvFilename(CSV_DIR, TOKEN_SYMBOL)
    if os.path.exists(c1):
        print(f"\nFile {c1} exists. Exiting.")
        sys.exit(0)
    c2 = csvs.rewardsinfoCsvFilename(CSV_DIR, TOKEN_SYMBOL)
    if os.path.exists(c2):
        print(f"\nFile {c2} exists. Exiting.")
        sys.exit(0)
    # main work
    stakes = csvs.loadStakesCsvs(CSV_DIR)
    poolvols = csvs.loadPoolvolsCsvs(CSV_DIR)
    rates = csvs.loadRateCsvs(CSV_DIR)
    rewardsperlp, rewardsinfo = calcrewards.calcRewards(
        stakes, poolvols, rates, TOT_TOKEN
    )
    csvs.saveRewardsperlpCsv(rewardsperlp, CSV_DIR, TOKEN_SYMBOL)
    csvs.saveRewardsinfoCsv(rewardsinfo, CSV_DIR, TOKEN_SYMBOL)

    print("dftool calc: Done")


# ========================================================================
@enforce_types
def do_dispense():
    HELP = f"""From rewards csv, dispense funds to chain.

Usage: dftool dispense CSV_DIR CHAINID [DFREWARDS_ADDR] [TOKEN_ADDR] [BATCH_NBR]
  CSV_DIR -- input directory for csv rewards file
  CHAINID -- DFRewards contract's network.{CHAINID_EXAMPLES}
  DFREWARDS_ADDR -- DFRewards contract's address. If not given, uses envvar DFREWARDS_ADDR
  TOKEN_ADDR -- token contract's address. If not given, uses envvar TOKEN_ADDR
  BATCH_NBR -- specify the batch number to run dispense only for that batch. If not given, runs dispense for all batches.

Transactions are signed with envvar 'DFTOOL_KEY`.
"""
    if len(sys.argv) not in [2 + 2, 2 + 3, 2 + 4]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "dispense"
    CSV_DIR = sys.argv[2]
    CHAINID = int(sys.argv[3])
    if len(sys.argv) == 4:
        print("Set DFREWARDS_ADDR from envvar")
        DFREWARDS_ADDR = os.getenv("DFREWARDS_ADDR")
    else:
        DFREWARDS_ADDR = sys.argv[4]

    if len(sys.argv) == 5:
        print("Set DFREWARDS_ADDR from envvar")
        TOKEN_ADDR = os.getenv("DFREWARDS_ADDR")
    else:
        TOKEN_ADDR = sys.argv[5]
    assert TOKEN_ADDR is not None

    batch_number = None
    if len(sys.argv) == 7:
        batch_number = sys.argv[6]

    print(
        f"Arguments: CSV_DIR={CSV_DIR}, CHAINID={CHAINID}"
        f", DFREWARDS_ADDR={DFREWARDS_ADDR}"
    )
    assert DFREWARDS_ADDR is not None

    # main work
    networkutil.connect(CHAINID)
    from_account = _getPrivateAccount()
    token_symbol = B.Simpletoken.at(TOKEN_ADDR).symbol()
    rewards = csvs.loadRewardsCsv(CSV_DIR, token_symbol)
    dispense.dispense(
        rewards[CHAINID],
        DFREWARDS_ADDR,
        TOKEN_ADDR,
        from_account,
        batch_number=batch_number,
    )

    print("dftool dispense: Done")


# ========================================================================
@enforce_types
def do_querymany():
    HELP = """Like 'query' except query *many* chains.

Usage: dftool querymany ST FIN NSAMP CSV_DIR CHAINID1 [CHAINID2 [CHAINID3 ..]]
"""
    if len(sys.argv) < 7:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "querymany"
    # pylint: disable=unbalanced-tuple-unpacking
    ST, FIN, NSAMP, CSV_DIR = sys.argv[2:6]
    CHAINIDS = [int(arg) for arg in sys.argv[6:]]
    print(f"dftool querymany: Begin. CHAINIDs={CHAINIDS}")

    for CHAINID in CHAINIDS:
        print(f"Call dftool query for CHAINID={CHAINID}: Begin")
        cmd = f"dftool query {ST} {FIN} {NSAMP} {CSV_DIR} {CHAINID}"
        os.system(cmd)
        print(f"Call dftool query for CHAINID={CHAINID}: Done")

    print("dftool querymany: Done")


# ========================================================================
@enforce_types
def do_newdfrewards():
    HELP = f"""Deploy new DFRewards contract.

Usage: dftool newdfrewards CHAINID
  CHAINID -- {CHAINID_EXAMPLES}
"""
    if len(sys.argv) not in [3]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "newdfrewards"
    CHAINID = int(sys.argv[2])

    print(f"Arguments: CHAINID={CHAINID}")

    # main work
    networkutil.connect(CHAINID)
    from_account = _getPrivateAccount()
    df_rewards = B.DFRewards.deploy({"from": from_account})
    print(f"New dispense contract deployed at address: {df_rewards.address}")

    print("dftool newdfrewards: Done")


# ========================================================================
@enforce_types
def do_compile():
    HELP = """Compile contracts

Usage: dftool compile
"""
    if len(sys.argv) not in [2]:
        print(HELP)
        sys.exit(0)

    os.system("brownie compile")


# ========================================================================
@enforce_types
def do_manyrandom():
    HELP = """Deploy many random tokens & pools then consume (for testing)

Usage: dftool manyrandom
"""
    if len(sys.argv) not in [2]:
        print(HELP)
        sys.exit(0)

    from util import oceantestutil  # pylint: disable=import-outside-toplevel

    recordDeployedContracts(oceantestutil._ADDRESS_FILE)
    oceantestutil.fillAccountsWithOCEAN()
    num_pools = 10  # magic number
    oceantestutil.randomDeployTokensAndPoolsThenConsume(num_pools, OCEANtoken())

    os.system("brownie compile")


# ========================================================================
@enforce_types
def do_mine():
    HELP = """Force chain to pass time (ganache only).

Usage: dftool mine BLOCKS [TIMEDELTA]
  BLOCKS -- e.g. 3
  TIMEDELTA -- e.g. 100
"""
    if len(sys.argv) not in [3, 4]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "mine"
    BLOCKS = int(sys.argv[2])
    if len(sys.argv) == 4:
        TIMEDELTA = int(sys.argv[3])
    else:
        TIMEDELTA = None

    print(f"Arguments: BLOCKS={BLOCKS}, TIMEDELTA={TIMEDELTA}")

    # main work
    networkutil.connect(DEV_CHAINID)
    chain = brownie.network.chain
    if TIMEDELTA is None:
        chain.mine(blocks=BLOCKS, timedelta=TIMEDELTA)
    else:
        chain.mine(blocks=BLOCKS)

    print("dftool mine: Done")


# ========================================================================
@enforce_types
def do_newacct():
    HELP = """Generate new account.

Usage: dftool newacct
"""
    if len(sys.argv) not in [2]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "newacct"

    # main work
    networkutil.connect(DEV_CHAINID)
    account = brownie.network.accounts.add()
    print("Generated new account:")
    print(f" private_key = {account.private_key}")
    print(f" address = {account.address}")
    print(f" For other dftools: export DFTOOL_KEY={account.private_key}")


# ========================================================================
@enforce_types
def do_newtoken():
    HELP = """Generate new token (for testing).

Usage: dftool newtoken CHAINID
"""
    if len(sys.argv) not in [3]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "newtoken"
    CHAINID = int(sys.argv[2])
    print(f"Arguments:\n CHAINID={CHAINID}")

    # main work
    networkutil.connect(CHAINID)
    from_account = _getPrivateAccount()
    token = B.Simpletoken.deploy("TST", "Test Token", 18, 1e21, {"from": from_account})
    print(f"Token '{token.symbol()}' deployed at address: {token.address}")


# ========================================================================
@enforce_types
def do_acctinfo():
    HELP = f"""Info about an account.

Usage: dftool acctinfo CHAINID ACCOUNT_ADDR [TOKEN_ADDR]
  CHAINID -- {CHAINID_EXAMPLES}
  ACCOUNT_ADDR -- e.g. '0x987...' or '4'. If the latter, uses accounts[i]
  TOKEN_ADDR -- e.g. '0x123..'

If envvar ADDRESS_FILE is not None, it gives balance for OCEAN token too.
"""
    if len(sys.argv) not in [4, 5]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "acctinfo"
    CHAINID = int(sys.argv[2])
    ACCOUNT_ADDR = sys.argv[3]
    TOKEN_ADDR = sys.argv[4] if len(sys.argv) >= 5 else None

    # do work
    print("Account info:")
    networkutil.connect(CHAINID)
    if len(str(ACCOUNT_ADDR)) == 1:
        addr_i = int(ACCOUNT_ADDR)
        ACCOUNT_ADDR = brownie.accounts[addr_i]
    print(f"  Address = {ACCOUNT_ADDR}")

    if TOKEN_ADDR is not None:
        token = B.Simpletoken.at(TOKEN_ADDR)
        balance = token.balanceOf(ACCOUNT_ADDR)
        print(f"  {fromBase18(balance)} {token.symbol()}")

    # Give balance for OCEAN token too.
    ADDRESS_FILE = os.environ.get("ADDRESS_FILE")
    if ADDRESS_FILE is not None:
        network = networkutil.chainIdToNetwork(CHAINID)
        recordDeployedContracts(ADDRESS_FILE, network)
        OCEAN = OCEANtoken()
        if OCEAN.address != TOKEN_ADDR:
            print(f"  {fromBase18(OCEAN.balanceOf(ACCOUNT_ADDR))} OCEAN")


# ========================================================================
@enforce_types
def do_chaininfo():
    HELP = f"""Info about a network.

Usage: dftool chaininfo CHAINID
  CHAINID -- {CHAINID_EXAMPLES}
"""
    if len(sys.argv) not in [3]:
        print(HELP)
        sys.exit(0)

    # extract inputs
    assert sys.argv[1] == "chaininfo"
    CHAINID = int(sys.argv[2])

    # do work
    networkutil.connect(CHAINID)
    # blocks = len(brownie.network.chain)
    print("\nChain info:")
    print(f"  # blocks: {len(brownie.network.chain)}")


# ========================================================================
# utilities


def _getAddressEnvvarOrExit() -> tuple:
    ADDRESS_FILE = os.environ.get("ADDRESS_FILE")
    print(f"Envvar:\n ADDRESS_FILE={ADDRESS_FILE}")
    if ADDRESS_FILE is None:
        print("\nNeed to set envvar ADDRESS_FILE. Exiting.")
        sys.exit(0)
    return ADDRESS_FILE


@enforce_types
def _getPrivateAccount():
    private_key = os.getenv("DFTOOL_KEY")
    assert private_key is not None, "Need to set envvar DFTOOL_KEY"
    account = brownie.network.accounts.add(private_key=private_key)
    print(f"For private key DFTOOL_KEY, address is: {account.address}")
    return account


# ========================================================================
# main
@enforce_types
def do_main():
    if len(sys.argv) == 1 or sys.argv[1] == "help":
        do_help()

    # write actions
    elif sys.argv[1] == "query":
        do_query()
    elif sys.argv[1] == "getrate":
        do_getrate()
    elif sys.argv[1] == "calc":
        do_calc()
    elif sys.argv[1] == "dispense":
        do_dispense()

    elif sys.argv[1] == "querymany":
        do_querymany()
    elif sys.argv[1] == "compile":
        do_compile()
    elif sys.argv[1] == "manyrandom":
        do_manyrandom()
    elif sys.argv[1] == "newdfrewards":
        do_newdfrewards()
    elif sys.argv[1] == "mine":
        do_mine()
    elif sys.argv[1] == "newacct":
        do_newacct()
    elif sys.argv[1] == "newtoken":
        do_newtoken()

    # read actions
    elif sys.argv[1] == "acctinfo":
        do_acctinfo()
    elif sys.argv[1] == "chaininfo":
        do_chaininfo()
    else:
        do_help()


if __name__ == "__main__":
    do_main()