ql-lg-server/tools/lgbot_setup.py

#!/usr/bin/env python3
"""LG training bot setup. Runs inside the container (through ./lgbot), like ./rcon.
  ./lgbot          questions: aim 1-5, movement 1-5, optional movement details; Enter keeps [value]
  ./lgbot --last   put the last setup back (e.g. after a server restart)
  ./lgbot show     saved setup and whether the bot is in
  ./lgbot off      remove the bot
Aim is the bot's skill level (addbot lgbot N). Movement is the lgbot_* cvars the lgbot plugin reads.
The last setup is saved in home/lgbot.json.
"""
import json, os, re, sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "home", "minqlx-plugins"))
from lgbot import PRESETS, LIMITS, KEYS  # noqa: E402  (the plugin's pattern settings)
from qlq import parse_rcon_status  # noqa: E402

STATE = os.path.join(HERE, "..", "home", "lgbot.json")
LABELS = {"strafe_min": "Strafe hold min ms", "strafe_max": "Strafe hold max ms",
          "fake": "Fake chance %", "range": "Preferred range (LG reach 768)",
          "range_play": "Range play: push in when ahead, back off when low y/n",
          "speed": "Strafe speed % of run speed",
          "lane": "Max sideways drift before turning back (no circling), 0 = off"}


def default_state():
    return {"aim": 3, "move": 3, "details": {}}


def load_state(path=STATE):
    st = default_state()
    try:
        with open(path) as f:
            raw = json.load(f)
    except (OSError, ValueError):
        return st
    if not isinstance(raw, dict):
        return st
    for k in ("aim", "move"):
        if isinstance(raw.get(k), int) and 1 <= raw[k] <= 5:
            st[k] = raw[k]
    details = raw.get("details")
    if isinstance(details, dict):
        for k, v in details.items():
            if k in LIMITS and isinstance(v, int) and LIMITS[k][0] <= v <= LIMITS[k][1]:
                st["details"][k] = v
    return st


def save_state(path, st):
    with open(path, "w") as f:
        json.dump(st, f, indent=1)


def parse_int(text, lo, hi, default):
    t = text.strip()
    if not t:
        return default
    try:
        v = int(t)
    except ValueError:
        return None
    return v if lo <= v <= hi else None


def parse_yn(text, default):
    t = text.strip().lower()
    if not t:
        return default
    return {"y": True, "yes": True, "n": False, "no": False}.get(t)


def ask(prompt, parse, input_fn=input, print_fn=print):
    while True:
        v = parse(input_fn(prompt))
        if v is not None:
            return v
        print_fn("  not valid, try again")


def prompt_setup(st, input_fn=input, print_fn=print):
    """Ask everything; Enter keeps the bracketed value. Ctrl-C/EOF propagate (caller cancels)."""
    aim = ask("Aim level 1-5 [%d]: " % st["aim"], lambda t: parse_int(t, 1, 5, st["aim"]), input_fn, print_fn)
    move = ask("Movement level 1-5 [%d]: " % st["move"], lambda t: parse_int(t, 1, 5, st["move"]), input_fn, print_fn)
    had = bool(st["details"])
    custom = ask("Customize movement details? %s: " % ("(Y/n)" if had else "(y/N)"),
                 lambda t: parse_yn(t, had), input_fn, print_fn)
    details = {}
    if custom:
        for k in KEYS:
            cur = st["details"].get(k, PRESETS[move][k])
            lo, hi = LIMITS[k]
            if k == "range_play":
                v = int(ask("  %s [%s]: " % (LABELS[k], "y" if cur else "n"),
                            lambda t, c=cur: parse_yn(t, bool(c)), input_fn, print_fn))
            else:
                v = ask("  %s %d-%d [%d]: " % (LABELS[k], lo, hi, cur),
                        lambda t, c=cur, lo=lo, hi=hi: parse_int(t, lo, hi, c), input_fn, print_fn)
            if v != PRESETS[move][k]:
                details[k] = v
    return {"aim": aim, "move": move, "details": details}


def _named_lgbot(status_text):
    parsed = parse_rcon_status(status_text)
    return [p for p in parsed[1] if p[2].lower() == "lgbot"] if parsed else []


def bot_present(status_text):
    return any(p[3] == "bot" for p in _named_lgbot(status_text))


def impostor(status_text):
    """A human calling themselves lgbot: "kick lgbot" would hit them, so don't send it."""
    return any(p[3] != "bot" for p in _named_lgbot(status_text))


def build_commands(st, readd, bot_in):
    cmds = ["set lgbot_move %d" % st["move"]]
    cmds += ["set lgbot_%s %d" % (k, st["details"].get(k, -1)) for k in KEYS]
    if readd:
        if bot_in:
            cmds.append("kick lgbot")
        cmds.append("addbot lgbot %d" % st["aim"])
    return cmds


def cvar_value(reply, name):
    m = re.search(r'"%s" is:"([^"]*)"' % re.escape(name), reply)
    return m.group(1) if m else None


def describe(st):
    d = ", ".join("%s %s" % (k, v) for k, v in sorted(st["details"].items()))
    return "aim %d, movement %d%s" % (st["aim"], st["move"], " (custom: %s)" % d if d else "")


def server():
    import qlrcon  # needs zmq + the RCON password: only available inside the container
    s = qlrcon.connect()
    return lambda cmd: qlrcon.run(s, cmd)


def refuse_impostor(status_text):
    if impostor(status_text):
        print("a player (not a bot) is named lgbot: rename or kick them first, nothing changed")
        return True
    return False


REFUSED = 3


def apply(st, old_aim):
    run = server()
    status = run("status")
    if refuse_impostor(status):
        return REFUSED
    bot_in = bot_present(status)
    for cmd in build_commands(st, readd=(not bot_in or st["aim"] != old_aim), bot_in=bot_in):
        run(cmd)
    ok = bot_present(run("status"))
    print("lgbot %s: %s" % ("in" if ok else "NOT in (is the server up? map loading?)", describe(st)))
    return 0 if ok else 1


def main(argv, input_fn=input):
    arg = argv[1] if len(argv) > 1 else ""
    st = load_state(STATE)
    if arg == "off":
        run = server()
        if refuse_impostor(run("status")):
            return 1
        run("kick lgbot")
        print("lgbot removed (setup kept, ./lgbot --last brings it back)")
        return 0
    if arg == "show":
        run = server()
        print("saved setup: " + describe(st))
        live = ["%s %s" % (k, cvar_value(run("lgbot_" + k), "lgbot_" + k)) for k in ["move"] + KEYS]
        print("live (-1 = from the movement level): " + ", ".join(live))
        print("bot is " + ("in" if bot_present(run("status")) else "not in"))
        return 0
    if arg == "--last":
        return apply(st, old_aim=None)  # old_aim None: always re-add so the aim level is right
    if arg:
        print(__doc__.strip())
        return 2
    try:
        new = prompt_setup(st, input_fn)
    except (EOFError, KeyboardInterrupt):
        print("\ncancelled, nothing changed")
        return 1
    rc = apply(new, old_aim=st["aim"])
    if rc != REFUSED:  # saved even if the bot didn't show up (server restarting): --last retries
        save_state(STATE, new)
    return rc


if __name__ == "__main__":
    sys.exit(main(sys.argv))