gear-orders/telegram/commands.py

import re
import logging
import peewee
import datetime

from db.queries import skip_day_put, skip_day_delete, skip_days_upcoming, user_add, user_get, domsubusers_add, domsubusers_delete
from .telegram import Telegram, TelegramCommand
from settings import FLASK_URL

logger = logging.getLogger(__name__)

class StartCommand(TelegramCommand):
  command_text = '/start'

  async def exec_inner(self, text, update, *args, **kwargs):
    username = update['message']['chat']['username']
    chat_id = update['message']['chat']['id']

    user_add(username, chat_id)

    yield "Welcome to Gear Orders bot.\n\nEdit your orders at {FLASK_URL}\n\nAdd a dom with `/dom_add <username>` or ask a sub to add you."

class DomAddCommand(TelegramCommand):
  command_regex = re.compile(r"^\/dom_add( (?P<username>@?\w+))$")

  async def exec_inner(self, text, update, session, *args, **kwargs):
    dom_username = text.split(' ')[1]
    if dom_username.startswith('@'):
      dom_username = dom_username[1:]

    sub_username = update['message']['chat']['username']
    sub = user_get(sub_username)
    try:
      dom = user_get(dom_username)
    except:
      yield "I'm sorry that user isn't registered with me."
      return
    
    try:
      domsubusers_add(sub, dom)
    except peewee.IntegrityError:
      yield f"@{dom_username} is already on your list of doms"
      return

    yield f"Successfully added {dom_username} to your list of doms"

    t = Telegram(session)
    await t.message_send(dom.telegram_chat_id, f"@{sub_username} has added you as a dom. You may edit their orders at {FLASK_URL}")

class DomRemoveCommand(TelegramCommand):
  command_regex = re.compile(r"^\/dom_remove( (?P<username>@?\w+))$")

  async def exec_inner(self, text, update, session, *args, **kwargs):
    dom_username = text.split(' ')[1]
    if dom_username.startswith('@'):
      dom_username = dom_username[1:]

    sub_username = update['message']['chat']['username']
    sub = user_get(sub_username)
    try:
      dom = user_get(dom_username)
    except:
      yield "I'm sorry that user isn't registered with me."
      return
    
    try:
      domsubusers_delete(sub, dom)
    except peewee.DoesNotExist:
      yield "I'm sorry that user is not on your list of doms"
      return

    yield f"Successfully removed {dom_username} from your list of doms"

    t = Telegram(session)
    await t.message_send(dom.telegram_chat_id, f"@{sub_username} has removed you as a dom.")


class SkipDayAddCommand(TelegramCommand):
  command_text = "/skip_day"

  async def exec_inner(self, text, update, session, forResponse=None, reply=None):
    try:
      yield "Please enter a skip day"

      response = await forResponse()

      user = user_get(update['message']['chat']['username'])

      skip_day = datetime.date.fromisoformat(response)
      skip_day_put(user, skip_day)

      yield f"Skip day {response} has been added"
    except ValueError:
      yield "That date was not valid"
    except peewee.IntegrityError:
      yield "That day has already been added"

class SkipDayDeleteCommand(TelegramCommand):
  command_text = "/skip_day_delete"

  async def exec_inner(self, text, update, session, forResponse=None, reply=None):
    try:
      yield "Please enter a skip day to delete"

      response = await forResponse()

      user = user_get(update['message']['chat']['username'])

      skip_day = datetime.date.fromisoformat(response)
      rows = skip_day_delete(user, skip_day)

      if rows > 0:
        yield f"Skip day {skip_day} has been deleted"
      else:
          yield f"Skip day {skip_day} does not exist"
    except ValueError:
      yield "That date was not valid"

class SkipDaysCommand(TelegramCommand):
  command_text = "/skip_days"

  async def exec_inner(self, text, update, session, forResponse=None, reply=None):
    user = user_get(update['message']['chat']['username'])
    dates = skip_days_upcoming(user)
    yield ("Upcoming skip days -\n" +
      "\n".join([d.date.isoformat() for d in dates]))

commands = [
  StartCommand(),
  DomAddCommand(),
  DomRemoveCommand(),
  SkipDayAddCommand(),
  SkipDayDeleteCommand(),
  SkipDaysCommand()
]