"""One-time setup for the starter bot, and its off switch.

    python setup_wallet.py            the whole setup: keys, funding, and in session mode the session
    python setup_wallet.py keys       make the keys (the wallet, and in session mode the executor) in key_dir
    python setup_wallet.py fund       say what to send where, and wait until it arrives (Ctrl+C stops waiting)
    python setup_wallet.py session    grant a session from config.toml's [session] caps, and send its setup
    python setup_wallet.py setup      send the session's setup again (after SESSION_CODE or SESSION_SETUP)
    python setup_wallet.py status     balances, the session, and what its caps have left
    python setup_wallet.py revoke     the off switch: end the session on chain, then at Orblivion
    python setup_wallet.py revoke --chain-only   end it on chain from session.json alone, without asking Orblivion
    python setup_wallet.py reset-code the emergency switch: make the wallet a plain account again (no smart-account code)

The keys are made on this computer, in key_dir (config.toml; default ~/.orblivion-bot), outside this folder, and are
never printed. Nothing is sent on chain unless you type SEND.

session.json keeps the session's revocation record (no secret: the delegation the wallet signed, its signature and what
the setup approved). With it, revoke builds the on-chain steps here and sends them straight to Robinhood Chain, so it
ends the session on chain even when Orblivion doesn't answer. That needs an SDK with Grant.revocation_record; with an
older one, revoke still works, with the on-chain steps from Orblivion's answer.
"""
import json
import os
import sys
import time

from bot import (EXPLORER, HERE, Ledger, client_for, key_path, load_account, load_config, load_session, read_book,
                 read_market, session_path)
from orblivion import OrblivionError

GAS_USD = 2.0         # what the executor should hold for gas: a session trade costs about 2 cents
MIN_GAS_USD = 1.0     # the executor counts as funded from here
MIN_WALLET_USD = 2.0  # the wallet counts as funded from here
MIN_WALLET_ETH_USD = 0.25   # the wallet's own ETH for gas (the setup, or every trade in direct mode)
UPDATE_SDK = "update the SDK for a revoke that needs nothing from Orblivion: pip install -U orblivion"


def confirm(*lines):
    """Nothing is sent unless SEND is typed."""
    print()
    for line in lines:
        print(line)
    if input("Type SEND to go ahead: ").strip() != "SEND":
        sys.exit("Not sent.")


def tx(h):
    return f"{EXPLORER}/tx/{h}"


def when(ts):
    return time.strftime("%Y-%m-%d %H:%M", time.localtime(ts))


def names(cfg):
    return ["wallet", "executor"] if cfg["mode"] == "session" else ["wallet"]


# --- keys -------------------------------------------------------------------------------------------------------------

def make_key(cfg, name):
    """A new key in <key_dir>/<name>.json, readable only by you where the system allows it. Never replaces a key."""
    from eth_account import Account
    os.makedirs(cfg["key_dir"], exist_ok=True)
    try:
        os.chmod(cfg["key_dir"], 0o700)
    except OSError:
        pass
    path = key_path(cfg, name)
    if os.path.exists(path):
        return False
    key = Account.create().key.hex()
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        json.dump({"private_key": key if key.startswith("0x") else "0x" + key}, f)
    return True


def keys(cfg):
    for name in names(cfg):
        made = make_key(cfg, name)
        print(f"{name:<9} {load_account(cfg, name).address}  ({'new' if made else 'already there'})")
    print(f"\nThe keys are in {cfg['key_dir']}. They never leave this computer and are never printed.")
    print("Anyone with wallet.json controls the wallet: never share or commit that folder, and back it up.")


# --- funding ----------------------------------------------------------------------------------------------------------

def fund(cfg, wait=True):
    wallet = load_account(cfg, "wallet")
    client = client_for(cfg, wallet)
    session = cfg["mode"] == "session"
    executor = load_account(cfg, "executor").address if session else None
    print(f"\nSend the wallet   {wallet.address}")
    print(f"  what the bot trades ({', '.join(cfg['tokens'])}), plus about $1 of ETH for "
          + ("its session setup." if session else "gas: each trade costs about a cent."))
    print("  Use a wallet that holds only what the bot trades.")
    if session:
        print(f"Send the executor {executor}")
        print(f"  about ${GAS_USD:g} of ETH: it pays the gas of every trade (about 2 cents each). It can't move the "
              "wallet's money.")
    print("Both are on Robinhood Chain (chain 4663).\n")
    while True:
        try:
            market = read_market(client, cfg["tokens"])
            book = read_book(client, cfg["tokens"], wallet.address, market)
            gas = read_book(client, ["ETH"], executor, market)["total_usd"] if session else None
        except OrblivionError as e:
            print(f"Couldn't read balances ({e.code}: {e.message}); trying again.")
            book = None
        except OSError as e:
            print(f"Couldn't read balances ({e}); trying again.")
            book = None
        if book is not None:
            held = ", ".join(f"{k} {v['units']:.6g}" for k, v in book["balances"].items() if v["units"] is not None)
            eth_usd = book["balances"].get("ETH", {}).get("usd") or 0
            line = f"wallet ${book['total_usd']:,.2f} ({held})"
            if session:
                line += f" | executor ${gas:,.2f} of ETH"
            print(f"{time.strftime('%H:%M:%S')}  {line}")
            gas_ok = not session or gas >= MIN_GAS_USD
            if book["total_usd"] >= MIN_WALLET_USD and eth_usd >= MIN_WALLET_ETH_USD and gas_ok:
                print("Funded.")
                return True
        if not wait:
            return False
        time.sleep(15)


# --- the session ------------------------------------------------------------------------------------------------------

def saved_session(cfg):
    path = session_path(cfg)
    if not os.path.exists(path):
        return None
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def save_session(cfg, s):
    with open(session_path(cfg), "w", encoding="utf-8") as f:
        json.dump(s, f, indent=1)


def revocation_record(grant):
    """The session's revocation record (Grant.revocation_record()), or None on an SDK without it."""
    fn = getattr(grant, "revocation_record", None)
    if not callable(fn):
        return None
    try:
        return fn()
    except (TypeError, ValueError):
        return None


def session(cfg):
    if cfg["mode"] != "session":
        sys.exit('config.toml says mode = "direct": there is no session to grant.')
    old = saved_session(cfg)
    if old and not old.get("revoked") and old.get("expires", 0) > time.time():
        sys.exit(f"Session {old['session_id']} is active until {when(old['expires'])}. To replace it, end it first: "
                 "python setup_wallet.py revoke")
    s, tokens = cfg["session"], cfg["tokens"]
    owner = client_for(cfg, load_account(cfg, "wallet"))
    executor = load_account(cfg, "executor").address
    budget = s["setup_budget_usd"] if s["setup_budget_usd"] is not None else s["max_day_usd"]
    print(f"Granting a session: the executor {executor} may trade {', '.join(tokens)} from the wallet "
          f"{owner.address}:")
    print(f"  ${s['max_trade_usd']:g} a trade, ${s['max_day_usd']:g} a day, {int(s['max_trades'])} trades, "
          f"for {s['hours']:g} hours; at most ${budget:g} of each ERC-20 it sells, over the whole session.")
    grant = owner.sessions.grant(
        executor=executor, tokens=tokens, max_trade_usd=s["max_trade_usd"], max_day_usd=s["max_day_usd"],
        max_trades=int(s["max_trades"]), hours=s["hours"], setup_budget_usd=budget,
        allow_7702=True,                      # your consent to set a plain wallet's smart-account code (see README)
        require_day_cap="ETH" in tokens,      # the chain caps ETH per 24 hours too
    )
    lim = grant.limits
    print(f"\nSession   {grant.session_id}, until {when(grant.expires)}")
    print(f"Caps      ${lim.max_trade_usd:g} a trade, ${lim.max_day_usd:g} a day, {lim.max_trades} trades, "
          f"{', '.join(lim.tokens)} only")
    print(f"Checked   {sum(c.ok for c in grant.checks)} of {len(grant.checks)} checks passed before signing")
    for x in grant.exposure or []:
        print(f"Note      the wallet has an approval outside the session; consider revoking it: {x}")
    record = revocation_record(grant)
    save_session(cfg, {"session_id": grant.session_id, "wallet": owner.address.lower(), "executor": executor.lower(),
                       "expires": grant.expires, "tokens": list(lim.tokens), "max_trade_usd": lim.max_trade_usd,
                       "max_day_usd": lim.max_day_usd, "max_trades": lim.max_trades, "setup_budget_usd": budget,
                       "setup": False, "revocation_record": record})
    print(f"Saved     {session_path(cfg)}" + (", with the session's revocation record" if record else ""))
    if record is None:
        print(f"Note      this SDK keeps no revocation record, so revoke will take its on-chain steps from Orblivion: "
              f"{UPDATE_SDK}")
    if grant.setup or grant.authorization_signed:
        confirm("The setup sends transactions from the wallet:",
                "  - on a plain wallet, one that sets its code to MetaMask's EIP7702StatelessDeleGator (EIP-7702), "
                "which sessions run on. The key stays the owner; an authorization to address 0 undoes it;",
                f"  - for each ERC-20 it sells, an approval of at most ${budget:g} through Permit2.")
        show_steps(owner.sessions.setup(grant))
    mark_setup(cfg)
    print("\nDone. Next: python bot.py (a dry run until you set dry_run = false in config.toml).")


def setup(cfg):
    s = load_session(cfg)
    owner = client_for(cfg, load_account(cfg, "wallet"))
    confirm(f"Send session {s['session_id']}'s setup from the wallet: its smart-account code if it's still a plain "
            "wallet, then any approval the session is missing.")
    d = owner.wallet.delegate_code()
    print(f"Code      {d.status}" + (f"  {tx(d.tx_hash)}" if d.tx_hash and d.status != "delegated" else ""))
    # each approval is held to what the wallet signed, and to the budget it was granted with
    show_steps(owner.sessions.setup_steps(s["session_id"], setup_budget_usd=s.get("setup_budget_usd")))
    mark_setup(cfg)


def show_steps(r):
    for step in r.steps:
        how = getattr(step, "status", None)
        print(f"  {step.kind}: {tx(step.tx_hash) if step.sent else 'not sent'}" + (f" ({how})" if how and how != "ok" else ""))
    if not r.steps:
        print("  nothing to send")


def mark_setup(cfg):
    s = saved_session(cfg)
    if s:
        s["setup"] = True
        save_session(cfg, s)


def status(cfg):
    owner = client_for(cfg, load_account(cfg, "wallet"))
    market = read_market(owner, cfg["tokens"])
    book = read_book(owner, cfg["tokens"], owner.address, market)
    print(f"Wallet    {owner.address}  ${book['total_usd']:,.2f}")
    for k, v in book["balances"].items():
        units = f"{v['units']:.8g}" if v["units"] is not None else "?"
        usd = f"${v['usd']:,.2f}" if v["usd"] is not None else "unpriced"
        print(f"          {k:<6} {units:<14} {usd}")
    for dry in (False, True):
        spent = Ledger(os.path.join(HERE, "bot-ledger.json"), time.time).spent(dry)
        kind = "dry runs" if dry else "live"
        print(f"Bot       ${spent:,.2f} of ${cfg['max_day_usd']:g} used in the last 24 h ({kind})")
    if cfg["mode"] != "session":
        return
    s = saved_session(cfg)
    if not s:
        print("Session   none yet: python setup_wallet.py session")
        return
    gas = read_book(owner, ["ETH"], s["executor"], market)["total_usd"]
    print(f"Executor  {s['executor']}  ${gas:,.2f} of ETH for gas" + ("" if gas >= MIN_GAS_USD else "  (top it up)"))
    print(f"Session   {s['session_id']}, until {when(s['expires'])}" + ("  REVOKED" if s.get("revoked") else ""))
    try:
        live = next((x for x in owner.me.sessions().sessions if str(x.session_id) == str(s["session_id"])), None)
    except OrblivionError as e:
        print(f"          couldn't read it from Orblivion: {e.code}: {e.message}")
        return
    if live is None:
        print("          not active at Orblivion (revoked or ended)")
        return
    print(f"          {live.status}; {live.trades_used} trades used, {live.trades_left} left; "
          f"${live.last_24h_usd or 0:,.2f} in the last 24 h, ${live.remaining_24h_usd or 0:,.2f} left today")


def revoke(cfg, chain_only=False):
    s = load_session(cfg)
    owner = client_for(cfg, load_account(cfg, "wallet"))
    record = s.get("revocation_record")
    local = bool(record) and hasattr(owner.sessions, "disable_onchain")
    if chain_only:
        return revoke_on_chain(cfg, s, owner, record, local)
    if local:
        how = ("They are built on this computer from session.json and sent straight to Robinhood Chain, even if Orblivion "
               "doesn't answer.")
    elif record:
        how = f"This SDK takes them from Orblivion's answer, so they need Orblivion to answer: {UPDATE_SDK}"
    else:
        how = ("session.json holds no revocation record (the session was granted with an older SDK), so they come from "
               "Orblivion's answer and need it to answer.")
    confirm(f"End session {s['session_id']} now: Orblivion stops co-signing at once, then transactions from the wallet "
            "end it on chain (the delegation disabled, the setup's approvals back to 0).", how)
    r = owner.sessions.revoke(s["session_id"], record=record) if local else owner.sessions.revoke(s["session_id"])
    api_error = getattr(r, "api_error", None)
    at_orblivion = not api_error and r.status == "revoked"
    print(f"Orblivion's own revoke failed ({api_error})" if api_error else f"Revoked at Orblivion: {r.effective or r.status}")
    show_steps(r)
    on = getattr(r, "on_chain", None)
    status = (on or {}).get("status")
    # everything the record holds reads done, and only Orblivion's answer (another root it names) stands in the way
    disputed = status == "unconfirmed" and (on or {}).get("from_record") == "confirmed"
    if status == "confirmed":
        print("On chain  confirmed: the delegation reads as disabled and the approvals as 0, on every RPC")
    elif status == "server_reported":
        print("On chain  sent, and read back as Orblivion named it: which roots were disabled is Orblivion's word "
              "(session.json holds no revocation record to check it)")
    elif status:
        print(f"On chain  {status}: " + "; ".join(on.get("detail") or []))
    ended = status == "confirmed" or (on is None and at_orblivion)       # an SDK without the read-back: as before
    if ended or at_orblivion or status == "server_reported":
        s["revoked"] = True
        save_session(cfg, s)
    if ended:
        print("The bot can't trade in it any more. For a new session: python setup_wallet.py session")
    elif status == "server_reported":
        print("Orblivion co-signs nothing more for it, so the bot can't trade in it, and the steps it named read back "
              "done. Running revoke again can't check it better: " + (
                  UPDATE_SDK if record else "session.json holds no revocation record (the session was granted with an "
                  "older SDK), so which roots ended is Orblivion's word. Sessions granted now keep one."))
    elif (on or {}).get("signed_live") and local:
        print(("Orblivion co-signs nothing more for it, so the bot can't trade in it. " if at_orblivion else "")
              + "On chain, Orblivion's answer carries a root your wallet signed that still reads live ("
              + ", ".join(h[:10] + "…" for h in on["signed_live"]) + "): session.json is most likely an older "
              "session's record, and that root this session's own.")
        disable_signed(cfg, s, owner, record)
    elif disputed:
        print(("Orblivion co-signs nothing more for it, so the bot can't trade in it. " if at_orblivion else "")
              + "On chain, everything session.json holds reads done on every RPC; only Orblivion's answer stands in the "
              "way (it names a root session.json doesn't hold, and none your wallet signed). Running revoke again won't "
              "change that. If session.json is this session's own record, take the chain's word from it alone: "
              "python setup_wallet.py revoke --chain-only")
    elif at_orblivion:
        print("Orblivion co-signs nothing more for it, so the bot can't trade in it; the chain doesn't confirm it ended "
              "yet: run python setup_wallet.py revoke again.")
    else:
        print("Neither Orblivion nor the chain confirms the session ended: run python setup_wallet.py revoke again.")


def disable_signed(cfg, s, owner, record):
    """The roots Orblivion's answer carries that the wallet signed and that read live, disabled too (behind SEND): the
    wallet's signature proves a root is the wallet's, not which of its sessions it is."""
    confirm("Disable that root too: one more transaction from the wallet, straight to Robinhood Chain.",
            "Your wallet signed it, so it is yours. If it is another session of yours that you still use, don't: this "
            "would end that one too.")
    r = owner.sessions.revoke(s["session_id"], record=record, disable_signed=True)
    show_steps(r)
    on = getattr(r, "on_chain", None) or {}
    if not on.get("signed_live") and on.get("from_record") == "confirmed":
        print("On chain  every root of session.json and every root your wallet signed that Orblivion named reads "
              "disabled, and the approvals 0, on every RPC")
        s["revoked"] = True
        save_session(cfg, s)
        print("The bot can't trade in it any more. For a new session: python setup_wallet.py session")
    else:
        print(f"On chain  {on.get('status')}: " + "; ".join(on.get("detail") or []))
        print("Not everything reads done yet: run python setup_wallet.py revoke again.")


def revoke_on_chain(cfg, s, owner, record, local):
    """revoke --chain-only: the on-chain steps from session.json's revocation record alone (disable_onchain), and the
    chain's answer read from it alone. Orblivion isn't asked anything."""
    if not local:
        sys.exit("revoke --chain-only ends the session from session.json's revocation record alone, so it needs " + (
            f"an SDK that can: {UPDATE_SDK}" if record else "a record, and session.json holds none (the session was "
            "granted with an older SDK): run python setup_wallet.py revoke"))
    confirm(f"End session {s['session_id']} on chain from session.json alone: transactions from the wallet disable its "
            "delegation and set the setup's approvals back to 0, sent straight to Robinhood Chain.",
            "Orblivion isn't asked anything, so it isn't told either: the disabled delegation refuses every trade on chain "
            "anyway.")
    r = owner.sessions.disable_onchain(record)
    show_steps(r)
    on = getattr(r, "on_chain", None) or {}
    if on.get("status") == "confirmed":
        print("On chain  confirmed from session.json alone: the delegation reads as disabled and the approvals as 0, on "
              "every RPC")
        s["revoked"] = True
        save_session(cfg, s)
        print("The bot can't trade in it any more. For a new session: python setup_wallet.py session")
    else:
        print(f"On chain  {on.get('status')}: " + "; ".join(on.get("detail") or []))
        print("The chain doesn't confirm the session ended yet: run python setup_wallet.py revoke --chain-only again.")


def reset_code(cfg):
    """The emergency switch: the wallet's own transaction to itself that sets its code back to none."""
    owner = client_for(cfg, load_account(cfg, "wallet"))
    st = owner.wallet.state()
    if not st.delegated_to:
        print(f"The wallet {owner.address} runs no smart-account code: there is nothing to reset.")
        return
    confirm("The emergency switch: make the wallet a plain account again. One transaction from the wallet to itself "
            "(an EIP-7702 authorization to address 0), sent straight to Robinhood Chain, not through Orblivion.",
            "What it does: the wallet runs no code any more, so no session can move anything from it. Every session "
            "trade fails on chain, whatever Orblivion or the executor do. Your key still controls the wallet, and your "
            "own transactions work as before.",
            "What it doesn't: the session isn't revoked (Orblivion still co-signs for it, and its delegation isn't "
            "disabled on chain), and the approvals stay (only the wallet's own transactions can use them). If the "
            "wallet's code is set again (python setup_wallet.py setup, or a new session), the old session works again "
            "until it expires: run python setup_wallet.py revoke as well.")
    r = owner.wallet.reset_code()
    print(f"Code      {r.status}" + (f"  {tx(r.tx_hash)}" if r.tx_hash else ""))
    if r.status == "reset":
        print("The wallet is a plain account again. To end the session as well: python setup_wallet.py revoke")
    elif r.status == "check":
        print(f"Sent, but the wallet's code still reads {r.code_after}: check the transaction before sending anything else")


def everything(cfg):
    keys(cfg)
    fund(cfg)
    if cfg["mode"] == "session":
        print()
        session(cfg)
    else:
        print("\nDone. Next: python bot.py (a dry run until you set dry_run = false in config.toml).")


def main(argv):
    cmd = argv[0] if argv else "all"
    acts = {"all": everything, "keys": keys, "fund": fund, "session": session, "setup": setup, "status": status,
            "revoke": lambda c: revoke(c, chain_only="--chain-only" in argv[1:]), "reset-code": reset_code}
    if cmd not in acts:
        sys.exit(__doc__)
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    cfg = load_config(os.path.join(HERE, "config.toml"))
    try:
        acts[cmd](cfg)
    except OrblivionError as e:
        sys.exit(f"\nRefused: {e.code}: {e.message}")
    except KeyboardInterrupt:
        sys.exit("\nStopped.")


if __name__ == "__main__":
    main(sys.argv[1:])