rollbot3/src/plugins/curse.py

from sqlalchemy import Column, Integer, String

from command_system import RollbotResponse, RollbotFailure, as_plugin, ModelBase, pop_arg


class CurseBlessScore(ModelBase):
    __tablename__ = "curse_bless_score"
    group_id = Column(String, primary_key=True)
    person_id = Column(Integer, primary_key=True)
    curses = Column(Integer)
    blessings = Column(Integer)


def get_score(db, group_id, name):
    person_id = name.strip().lower()
    score = db.query(CurseBlessScore).get((group_id, person_id))
    if not score:
        score = CurseBlessScore(group_id=group_id, person_id=person_id, curses=0, blessings=0)
        db.add(score)
    return score


def fmt_times(n):
    return str(n) + (" time" if n == 1 else " times")


def get_response(msg, name, score):
    return RollbotResponse(msg, txt=f"It is done! {name} has been cursed {fmt_times(score.curses)} and blessed {fmt_times(score.blessings)} in this chat.")


@as_plugin
def curse(db, msg):
    # TODO might be nice to add subcommands to this later
    name, args = pop_arg(msg.raw_args)
    if not name.startswith("!"):
        score = get_score(db, msg.group_id, name)
        # Note we do this instead of += b/c that can create race conditions in sqlalchemy
        score.curses = score.curses + 1
        return get_response(msg, name, score)

    # strip off the '!'
    subc = name[1:].strip()
    if len(subc) == 0:
        # handle the case of spaces between ! and subcommand
        subc, args = pop_arg(args)
    subc = subc.lower()
    name, _ = pop_arg(args)
    person_id = name.strip().lower()

    if subc == "aggregate":
        curses = 0
        blessings = 0
        for cbs in db.query(CurseBlessScore).filter(CurseBlessScore.person_id == person_id):
            curses += cbs.curses
            blessings += cbs.blessings
        return RollbotResponse(msg, txt=f"From my records, {name} has been cursed {fmt_times(curses)} and blessed {fmt_times(blessings)} across all chats.")
    elif subc == "clear":
        if not msg.from_admin:
            return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS)
        cbs = db.query(CurseBlessScore).get((msg.group_id, person_id))
        if cbs:
            oc = cbs.curses
            ob = cbs.blessings
            cbs.curses = 0
            cbs.blessings = 0
            return RollbotResponse(msg, txt=f"Done! I have cleared all records of {name} in this chat, who was previously cursed {fmt_times(oc)} and blessed {fmt_times(ob)}!")
        return RollbotResponse(msg, txt=f"Sorry! I don't know who {name} is!")
    elif subc == "clearall":
        if not msg.from_admin:
            return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS)
        cnt = 0
        oc = 0
        ob = 0
        for cbs in db.query(CurseBlessScore).filter(CurseBlessScore.person_id == person_id):
            oc += cbs.curses
            ob += cbs.blessings
            cbs.curses = 0
            cbs.blessings = 0
            cnt += 1
        if cnt == 0:
            return RollbotResponse(msg, txt=f"Sorry! I don't know who {name} is!")
        return RollbotResponse(msg, txt=f"Done! I have cleared all records of {name} in all chats, who was previously cursed {fmt_times(oc)} and blessed {fmt_times(ob)} in total!")
    else:
        return RollbotResponse(msg, failure=RollbotFailure.INVALID_SUBCOMMAND)



@as_plugin
def bless(db, msg):
    name, _ = pop_arg(msg.raw_args)
    score = get_score(db, msg.group_id, name)
    score.blessings = score.blessings + 1
    return get_response(msg, name, score)