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 db.models import OrdersPool
from .telegram import Telegram, TelegramCommand
from settings import FLASK_URL

logger = logging.getLogger(__name__)

CALLBACK_PREFIX_SKIP_DAY_ADD = "skip_day_add_"

class StartCommand(TelegramCommand):
  command_text = '/start'
  description = "Start a session with the bot"

  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 f"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_text = "/dom_add"
  command_regex = re.compile(r"^\/dom_add( (?P<username>@?\w+))$")
  description = "Add a dom to your profile"

  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_text = "/dom_remove"
  command_regex = re.compile(r"^\/dom_remove( (?P<username>@?\w+))$")
  description = "Remove a dom from your profile"

  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"
  description = "Add a skip day"

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

      skip_day_str = await forResponse()

      try:
        skip_day = datetime.date.fromisoformat(skip_day_str)
      except:
        yield "The date you entered was not valid"
        return

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

      inline_keyboard = [{
        'text': 'All',
        'callback_data': f"{CALLBACK_PREFIX_SKIP_DAY_ADD}_all"
      }] + [
        {
          'text': pool.name,
          'callback_data': f"{CALLBACK_PREFIX_SKIP_DAY_ADD}_{pool.id}"
        } for pool in user.orders_pools
      ]

      yield "Please choose an order pool", {
        "inline_keyboard": [inline_keyboard]
      }

      pool_str = await forResponse()

      if not pool_str.startswith("skip_day_add_"):
        yield "There was a problem with your command"
        return

      try:
        order_pool_id = int(pool_str[len(CALLBACK_PREFIX_SKIP_DAY_ADD)+1:])
        order_pool = user.orders_pools.where(OrdersPool.id == order_pool_id).get()
      except:
        order_pool = None

      skip_day_put(user, skip_day, order_pool)

      if order_pool is not None:
        yield f"Skip day {skip_day_str} for pool {order_pool.name} has been added"
      else:
        yield f"Skip day {skip_day_str} 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"
  description = "Delete a skip day"

  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"
  description = "List upcoming 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)
    response = "Upcoming skip days -\n"

    for date in dates:
      if date.pool is None:
        response += date.date.isoformat() + "\n"
      else:
        response += f"{date.date.isoformat()} - {date.pool.name}\n"

    yield response

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