"""The starter bot's loop. Your logic goes in strategy.py and your settings in config.toml: you shouldn't need to edit
this file.

Every interval_s seconds it:
  1. reads the USD prices of the configured tokens (Chainlink, read on chain by the SDK) and the wallet's balances;
  2. asks strategy.decide(market, book, state) what to trade;
  3. cuts each order to the bot's own limits (per trade, per rolling 24 hours, what the wallet holds, at least $1);
  4. trades it, or with dry_run on (the default) only checks it, and logs everything to the console and bot-log.jsonl.

    python bot.py                  run until Ctrl+C, or until a file named STOP appears in this folder
    python bot.py --once           one tick, then exit
    python bot.py --config my.toml

A dry run never uses the session's limits. Orblivion counts every trade its co-signer signs against the session's day
cap, sent or not, so a dry run here never asks for a session trade: it asks for a quote (session mode) or a plan (direct
mode), and nothing is co-signed, signed or sent.
"""
import argparse
import importlib.util
import json
import math
import os
import sys
import time

try:
    from orblivion import OrblivionApiError, OrblivionError, VerificationError
except ImportError:  # pragma: no cover
    sys.exit('The Orblivion SDK is missing: pip install "orblivion>=0.1.1"')

try:
    import tomllib  # Python 3.11 and later
except ModuleNotFoundError:  # Python 3.10: the small reader below
    tomllib = None

HERE = os.path.dirname(os.path.abspath(__file__))
EXPLORER = "https://robin.etherscan.io"
MIN_TRADE_USD = 1.0
DAY_S = 24 * 3600
KEEP_BACK = 0.98          # an order spends at most 98% of what the wallet holds of a token (gas, rounding)

# Orblivion's refusals after which this session can't trade again: the bot stops and says what to do.
SESSION_OVER = {
    "SESSION_REVOKED": "the session was revoked",
    "SESSION_EXPIRED": "the session has ended",
    "SESSION_EXPIRY": "the session is about to end",
    "SESSION_PENDING": "the session isn't switched on",
    "SESSION_TRADES_LEFT": "the session's trades are used up",
    "SESSION_CODE": "the wallet can't run the session's code yet",
}
# What to do about some refusals the bot can carry on after.
HINTS = {
    "SESSION_SETUP": "the wallet hasn't approved this token for the session: python setup_wallet.py setup",
    "SESSION_MAX_DAY": "the session's day cap is used up; it frees up as the 24 hours roll on",
    "MAX_NOTIONAL": "over the SDK's backstop, a little above limits.max_trade_usd",
    "INSUFFICIENT_FUNDS": "send the wallet (or in session mode the executor, for gas) more ETH",
}


# --- settings ---------------------------------------------------------------------------------------------------------

def _strip_comment(line):
    out, quote = [], None
    for ch in line:
        if quote:
            quote = None if ch == quote else quote
        elif ch in "\"'":
            quote = ch
        elif ch == "#":
            break
        out.append(ch)
    return "".join(out).strip()


def _split_items(s):
    items, cur, quote = [], "", None
    for ch in s:
        if quote:
            quote = None if ch == quote else quote
        elif ch in "\"'":
            quote = ch
        elif ch == ",":
            items.append(cur)
            cur = ""
            continue
        cur += ch
    if cur.strip():
        items.append(cur)
    return items


def _toml_value(raw, where):
    v = raw.strip()
    if v.startswith("[") and v.endswith("]"):
        return [_toml_value(x, where) for x in _split_items(v[1:-1]) if x.strip()]
    if len(v) >= 2 and v[0] == v[-1] and v[0] in "\"'":
        return v[1:-1]
    if v in ("true", "false"):
        return v == "true"
    for kind in (int, float):
        try:
            return kind(v.replace("_", ""))
        except ValueError:
            pass
    raise ValueError(f"{where}: can't read {raw.strip()!r}. On Python 3.10 this bot reads simple TOML only (strings, "
                     "numbers, true/false, one-line lists); Python 3.11 or later reads any TOML.")


def read_toml(text):
    """The part of TOML config.toml uses, for Python 3.10 (3.11 and later use the standard library's tomllib)."""
    data, table = {}, None
    for n, line in enumerate(text.splitlines(), 1):
        line = _strip_comment(line)
        if not line:
            continue
        if line.startswith("[") and line.endswith("]"):
            table = data.setdefault(line[1:-1].strip(), {})
            continue
        key, eq, value = line.partition("=")
        if not eq:
            raise ValueError(f"line {n}: expected key = value")
        (data if table is None else table)[key.strip()] = _toml_value(value, f"line {n}")
    return data


def load_config(path):
    """config.toml, checked: a setting that doesn't make sense stops the bot before it does anything."""
    with open(path, encoding="utf-8") as f:
        text = f.read()
    data = tomllib.loads(text) if tomllib else read_toml(text)
    bot, limits, session = data.get("bot", {}), data.get("limits", {}), data.get("session", {})
    cfg = {
        "tokens": bot.get("tokens", ["ETH", "USDG"]),
        "interval_s": bot.get("interval_s", 60),
        "dry_run": bot.get("dry_run", True),
        "mode": bot.get("mode", "session"),
        "key_dir": os.path.abspath(os.path.expanduser(bot.get("key_dir", "~/.orblivion-bot"))),
        "api_url": bot.get("api_url") or None,      # None: ORBLIVION_API_URL, else https://api.orblivion.com
        "stop_file": bot.get("stop_file", "STOP"),
        "max_trade_usd": limits.get("max_trade_usd", 10),
        "max_day_usd": limits.get("max_day_usd", 40),
        "session": {
            "max_trade_usd": session.get("max_trade_usd", 10),
            "max_day_usd": session.get("max_day_usd", 40),
            "max_trades": session.get("max_trades", 200),
            "hours": session.get("hours", 168),
            "setup_budget_usd": session.get("setup_budget_usd"),
        },
    }
    problems = []
    if not isinstance(cfg["dry_run"], bool):
        problems.append("dry_run is true or false (no quotes)")
    if cfg["mode"] not in ("session", "direct"):
        problems.append('mode is "session" or "direct"')
    toks = cfg["tokens"]
    if not (isinstance(toks, list) and len(toks) >= 2 and all(isinstance(t, str) and t for t in toks)
            and len(set(toks)) == len(toks)):
        problems.append('tokens is a list of at least two different tokens, e.g. ["ETH", "USDG"]')
    if not _num(cfg["interval_s"]) or cfg["interval_s"] < 5:
        problems.append("interval_s is a number of seconds, at least 5")
    if not _num(cfg["max_trade_usd"]) or cfg["max_trade_usd"] < MIN_TRADE_USD:
        problems.append("limits.max_trade_usd is a number of dollars, at least 1")
    if not _num(cfg["max_day_usd"]) or cfg["max_day_usd"] < cfg["max_trade_usd"]:
        problems.append("limits.max_day_usd is at least limits.max_trade_usd")
    s = cfg["session"]
    if not all(_num(s[k]) for k in ("max_trade_usd", "max_day_usd", "max_trades", "hours")):
        problems.append("session.max_trade_usd, max_day_usd, max_trades and hours are numbers")
    elif not 0 < s["hours"] <= 720:
        problems.append("session.hours is above 0 and at most 720 (30 days)")
    if problems:
        raise SystemExit(f"{path}:\n  " + "\n  ".join(problems))
    return cfg


def _num(x):
    return isinstance(x, (int, float)) and not isinstance(x, bool) and math.isfinite(x)


# --- keys and the session file (in key_dir, outside this folder) ------------------------------------------------------

def key_path(cfg, name):
    return os.path.join(cfg["key_dir"], name + ".json")


def load_account(cfg, name):
    """The key in <key_dir>/<name>.json, as an eth_account account. The key itself is never printed or logged."""
    from eth_account import Account
    path = key_path(cfg, name)
    if not os.path.exists(path):
        raise SystemExit(f"No {name} key in {cfg['key_dir']}: run python setup_wallet.py first.")
    with open(path, encoding="utf-8") as f:
        return Account.from_key(json.load(f)["private_key"])


def session_path(cfg):
    return os.path.join(cfg["key_dir"], "session.json")


def load_session(cfg):
    path = session_path(cfg)
    if not os.path.exists(path):
        raise SystemExit(f"No session in {cfg['key_dir']}: run python setup_wallet.py (session mode needs one).")
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def client_for(cfg, account):
    """An Orblivion client for this key. max_notional_usd is the SDK's own backstop, a little above the bot's per-trade
    limit so that a price move between the order and the SDK's check doesn't refuse an order the bot already capped."""
    from orblivion import OrblivionClient
    return OrblivionClient(base_url=cfg["api_url"], account=account, max_notional_usd=cfg["max_trade_usd"] * 1.1)


# --- reading prices and balances (no key needed: public chain reads) --------------------------------------------------

def read_market(client, tokens):
    """USD prices from Chainlink, read on chain by the SDK. USDG is $1; a token without a feed is left out."""
    marks = client.marks(list(tokens))
    market = {sym: m.price8 / 1e8 for sym, m in marks.items() if sym in tokens}
    if "WETH" in tokens and "ETH" in marks:
        market["WETH"] = marks["ETH"].price8 / 1e8
    return market


def read_book(client, tokens, who, market, decimals=None):
    """``who``'s balances, read on chain: units, USD (None without a price) and the priced total. Each token's decimals
    go into ``decimals`` when given."""
    names = {t.lower() if t.startswith("0x") else t: t for t in tokens}
    balances, total = {}, 0.0
    for b in client.wallet.balances(list(tokens), who=who):
        name = names.get(b.token, b.token)
        units = None if b.balance is None or b.decimals is None else b.balance / 10 ** b.decimals
        if decimals is not None and b.decimals is not None:
            decimals[name] = b.decimals
        price = market.get(name)
        usd = None if units is None or price is None else units * price
        balances[name] = {"units": units, "usd": usd}
        total += usd or 0
    return {"balances": balances, "total_usd": total}


# --- the bot's own records --------------------------------------------------------------------------------------------

class Ledger:
    """What the bot has spent in the last 24 hours, in a JSON file next to it. Dry runs and live trades are counted
    apart, so a dry run shows what the limits would let through, and going live starts with a clean day."""

    def __init__(self, path, clock):
        self.path, self.clock = path, clock
        try:
            with open(path, encoding="utf-8") as f:
                self.rows = json.load(f)["rows"]
        except (OSError, ValueError, KeyError):
            self.rows = []

    def spent(self, dry):
        since = self.clock() - DAY_S
        return sum(r["usd"] for r in self.rows if r["ts"] > since and r["dry"] == dry)

    def add(self, usd, dry, sell, buy):
        row = {"ts": self.clock(), "usd": usd, "dry": dry, "sell": sell, "buy": buy}
        self.rows.append(row)
        self.save()
        return row

    def remove(self, row):
        if row in self.rows:
            self.rows.remove(row)
            self.save()

    def save(self):
        since = self.clock() - DAY_S
        self.rows = [r for r in self.rows if r["ts"] > since]
        with open(self.path, "w", encoding="utf-8") as f:
            json.dump({"rows": self.rows}, f, indent=1)


class Log:
    """Each line to the console, and as JSON to bot-log.jsonl."""

    def __init__(self, path, clock, out=print):
        self.path, self.clock, self.out = path, clock, out

    def __call__(self, kind, msg, **data):
        ts = self.clock()
        with open(self.path, "a", encoding="utf-8") as f:
            f.write(json.dumps(dict({"ts": round(ts, 3), "kind": kind, "msg": msg}, **data), default=str) + "\n")
        self.out(f"{time.strftime('%H:%M:%S', time.localtime(ts))}  {kind:<6} {msg}")


# --- the bot ----------------------------------------------------------------------------------------------------------

class Bot:
    """``client`` is an OrblivionClient: the executor's in session mode, the wallet's in direct mode. ``wallet`` is the
    wallet that holds the money; ``session_id`` the session the executor trades in (session mode)."""

    def __init__(self, cfg, client, wallet, session_id=None, decide=None, folder=HERE, clock=time.time,
                 sleep=time.sleep, out=print):
        if decide is None:
            decide = load_strategy().decide
        self.cfg, self.client, self.wallet, self.session_id, self.decide = cfg, client, wallet, session_id, decide
        self.folder, self.clock, self.sleep = folder, clock, sleep
        self.ledger = Ledger(os.path.join(folder, "bot-ledger.json"), clock)
        self.log = Log(os.path.join(folder, "bot-log.jsonl"), clock, out)
        self.state_path = os.path.join(folder, "bot-state.json")
        self.state = self._load_state()
        self.decimals = {}
        self.wait_until = 0.0     # set when Orblivion asks the bot to wait before trying again
        self.stopped = None       # why the bot must stop (a session that can't trade any more)
        self._state_warned = False

    @property
    def session_mode(self):
        return self.cfg["mode"] == "session"

    # -- reading -------------------------------------------------------------------------------------------------------

    def read_market(self):
        return read_market(self.client, self.cfg["tokens"])

    def read_book(self, market):
        return read_book(self.client, self.cfg["tokens"], self.wallet, market, self.decimals)

    # -- one tick ------------------------------------------------------------------------------------------------------

    def tick(self):
        if self.clock() < self.wait_until:
            self.log("wait", f"waiting {self.wait_until - self.clock():.0f} s more, as Orblivion asked")
            return
        try:
            market = self.read_market()
            book = self.read_book(market)
        except Exception as e:  # noqa: BLE001  (an RPC or network failure: try again next tick)
            self.error("reading prices and balances", e)
            return
        prices = ", ".join(f"{k} ${v:,.2f}" for k, v in market.items())
        held = ", ".join(f"{k} {v['units']:.6g}" for k, v in book["balances"].items() if v["units"] is not None)
        self.log("tick", f"{prices or 'no prices'} | {held or 'no balances'} | ${book['total_usd']:,.2f}",
                 market=market, book=book)
        try:
            orders = self.decide(market, book, self.state) or []
            if not isinstance(orders, (list, tuple)):
                orders = [orders]
        except Exception as e:  # noqa: BLE001  (a bug in strategy.py mustn't stop the bot)
            self.log("error", f"strategy.decide raised {type(e).__name__}: {e}")
            orders = []
        self._save_state()
        if not orders:
            self.log("hold", "no orders")
        for order in orders:
            if self.stopped:
                break
            self.place(order, book)

    def check(self, order):
        """An order as (sell, buy, amount_usd, reason), or None (logged) when it isn't one the bot can place."""
        get = order.get if isinstance(order, dict) else lambda k, d=None: getattr(order, k, d)
        sell, buy, usd, reason = get("sell"), get("buy"), get("amount_usd"), get("reason", "") or ""
        if sell not in self.cfg["tokens"] or buy not in self.cfg["tokens"] or sell == buy:
            self.log("skip", f"{order!r}: sell and buy must be two different tokens from config.toml's tokens")
            return None
        if not _num(usd) or usd <= 0:
            self.log("skip", f"{order!r}: amount_usd must be a number above 0")
            return None
        return sell, buy, float(usd), str(reason)

    def place(self, order, book):
        checked = self.check(order)
        if checked is None:
            return
        sell, buy, want, reason = checked
        dry = self.cfg["dry_run"]
        left_today = self.cfg["max_day_usd"] - self.ledger.spent(dry)
        have = (book["balances"].get(sell) or {}).get("usd")
        limits = {"max_trade_usd": self.cfg["max_trade_usd"], "max_day_usd left": left_today}
        if have is not None:
            limits["the wallet's " + sell] = have * KEEP_BACK
        usd = math.floor(min(want, *limits.values()) * 100) / 100
        cut = [k for k, v in limits.items() if v < want and math.floor(v * 100) / 100 == usd]
        if usd < MIN_TRADE_USD:
            self.log("skip", f"{sell} -> {buy}: ${max(usd, 0):,.2f} after the limits ({', '.join(cut) or 'too small'}),"
                             f" under ${MIN_TRADE_USD:g}; asked ${want:,.2f}", reason=reason)
            return
        note = f" (asked ${want:,.2f}; cut to {', '.join(cut)})" if cut else ""
        self.log("order", f"sell ${usd:,.2f} of {sell} for {buy}{note}" + (f": {reason}" if reason else ""),
                 sell=sell, buy=buy, usd=usd,
                 asked_usd=want, reason=reason, dry_run=dry)
        row = self.ledger.add(usd, dry, sell, buy)
        try:
            if dry:
                self.dry_run(sell, buy, usd)
            else:
                self.trade(sell, buy, usd)
        except Exception as e:  # noqa: BLE001
            if dry or _nothing_sent(e):
                self.ledger.remove(row)     # nothing moved: the day's limit gets it back
            self.error(f"{'dry run' if dry else 'trade'} {sell} -> {buy}", e)

    def dry_run(self, sell, buy, usd):
        if self.session_mode:
            # A quote: read-only. Not sessions.trade(dry_run=True), which has the co-signer sign a leaf that counts
            # against the session's day cap even though it is never sent.
            q = self.client.quote(sell=sell, buy=buy, amount_usd=usd)
            out = q.best.amount_out if q.best else None
            fee = "?" if q.fee_bps is None else q.fee_bps
            self.log("dry", f"would sell {self.units(q.amount_in, sell)} {sell} for about {self.units(out, buy)} {buy}"
                            f" ({fee} bps fee); a quote: nothing co-signed or sent",
                     sell=sell, buy=buy, usd=usd, amount_in=q.amount_in, amount_out=out, fee_bps=q.fee_bps)
        else:
            # A plan: every check the real trade gets, with nothing signed or sent.
            p = self.client.plan(sell=sell, buy=buy, amount_usd=usd)
            passed = sum(1 for c in p.checks if c.ok)
            self.log("dry", f"would sell {self.units(p.amount_in, sell)} {sell} for about "
                            f"{self.units(p.simulated_out, buy)} {buy}, at least {self.units(p.min_out, buy)}; "
                            f"{passed} of {len(p.checks)} checks passed; nothing signed or sent",
                     sell=sell, buy=buy, usd=usd, amount_in=p.amount_in, simulated_out=p.simulated_out,
                     min_out=p.min_out, fee_bps=p.fee_bps)

    def trade(self, sell, buy, usd):
        if self.session_mode:
            t = self.client.sessions.trade(self.session_id, sell=sell, buy=buy, amount_usd=usd, wallet=self.wallet)
        else:
            t = self.client.trade(sell=sell, buy=buy, amount_usd=usd)
        self.log("trade", f"sold {self.units(t.amount_in, sell)} {sell} for {self.units(t.amount_out, buy)} {buy}  "
                          f"{EXPLORER}/tx/{t.tx_hash}", sell=sell, buy=buy, usd=usd, tx_hash=t.tx_hash,
                 amount_in=t.amount_in, amount_out=t.amount_out, spend=getattr(t, "spend", None))

    def error(self, what, e):
        if isinstance(e, OrblivionError):
            hint = HINTS.get(e.code)
            self.log("error", f"{what}: {e.code}: {e.message}" + (f" ({hint})" if hint else ""), code=e.code,
                     retryable=e.retryable, retry_after=e.retry_after)
            if e.code in SESSION_OVER and self.session_mode:
                fix = ("python setup_wallet.py setup" if e.code == "SESSION_CODE"
                       else "python setup_wallet.py session, for a new session")
                self.stopped = f"{SESSION_OVER[e.code]} ({e.code}). Run {fix}."
            elif e.retryable:
                wait = e.retry_after if e.retry_after else max(self.cfg["interval_s"], 30)
                self.wait_until = self.clock() + wait
        else:
            self.log("error", f"{what}: {type(e).__name__}: {e} (trying again next tick)")

    # -- the loop ------------------------------------------------------------------------------------------------------

    def stop_requested(self):
        return os.path.exists(os.path.join(self.folder, self.cfg["stop_file"]))

    def run(self, once=False):
        c = self.cfg
        how = "DRY RUN: nothing is sent" if c["dry_run"] else "LIVE: trades are sent"
        self.log("start", f"{c['mode']} mode, {how}; wallet {self.wallet}; every {c['interval_s']:g} s; at most "
                          f"${c['max_trade_usd']:g} a trade and ${c['max_day_usd']:g} a day")
        try:
            while not self.stop_requested():
                self.tick()
                if once or self.stopped:
                    break
                self.nap(c["interval_s"])
            else:
                self.log("stop", f"found {c['stop_file']}: stopping (delete it to run again)")
        except KeyboardInterrupt:
            self.log("stop", "Ctrl+C: stopped")
        finally:
            self._save_state()
        if self.stopped:
            self.log("stop", self.stopped)
            return 1
        return 0

    def nap(self, seconds):
        end = self.clock() + seconds
        while self.clock() < end and not self.stop_requested():
            self.sleep(min(1.0, end - self.clock()))

    # -- helpers -------------------------------------------------------------------------------------------------------

    def units(self, atoms, token):
        d = self.decimals.get(token)
        if atoms is None or d is None:
            return "?"
        return f"{atoms / 10 ** d:.6g}"

    def _load_state(self):
        try:
            with open(self.state_path, encoding="utf-8") as f:
                state = json.load(f)
            return state if isinstance(state, dict) else {}
        except (OSError, ValueError):
            return {}

    def _save_state(self):
        try:
            text = json.dumps(self.state, indent=1)
        except (TypeError, ValueError) as e:
            if not self._state_warned:
                self.log("error", f"state can't be saved to bot-state.json ({e}): keep plain values in it")
                self._state_warned = True
            return
        with open(self.state_path, "w", encoding="utf-8") as f:
            f.write(text)


def load_strategy():
    """strategy.py, from this bot's own folder."""
    spec = importlib.util.spec_from_file_location("strategy", os.path.join(HERE, "strategy.py"))
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def _nothing_sent(e):
    """True when the error proves nothing was sent: a local check refused to sign, or Orblivion refused the request."""
    return isinstance(e, VerificationError) or (isinstance(e, OrblivionApiError) and 400 <= e.status < 500)


def main(argv=None):
    ap = argparse.ArgumentParser(description="The Orblivion starter bot: trades what strategy.py decides.")
    ap.add_argument("--config", default=os.path.join(HERE, "config.toml"))
    ap.add_argument("--once", action="store_true", help="run one tick, then exit")
    args = ap.parse_args(argv)
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    cfg = load_config(args.config)
    if cfg["mode"] == "session":
        s = load_session(cfg)
        if s.get("revoked"):
            raise SystemExit("The saved session was revoked: run python setup_wallet.py session for a new one.")
        if s.get("expires", 0) <= time.time():
            raise SystemExit("The saved session has ended: run python setup_wallet.py session for a new one.")
        if sorted(s.get("tokens", [])) != sorted(cfg["tokens"]):
            print(f"Note: the session trades {', '.join(s.get('tokens', []))}; config.toml lists "
                  f"{', '.join(cfg['tokens'])}. Orblivion refuses trades outside the session.")
        bot = Bot(cfg, client_for(cfg, load_account(cfg, "executor")), wallet=s["wallet"], session_id=s["session_id"])
    else:
        client = client_for(cfg, load_account(cfg, "wallet"))
        bot = Bot(cfg, client, wallet=client.address)
    return bot.run(once=args.once)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(0)