path: root/velascobot2/bot.py

import random
import logging
import time
import sys

from telegram import Chat, ChatMember, Message, Update
import telegram
from telegram.ext import (
    Application,
    CommandHandler,
    ContextTypes,
    MessageHandler,
    filters,
)
from telegram.error import NetworkError

from velascobot2.utils import ChatData, GraphicType, Response
from velascobot2.messenger import Messenger

logger = logging.getLogger(__name__)


start_msg = (
    "Hello there! Ask me for /help to see an overview of the available"
    " commands."
)


wake_msg = "Good morning. I just woke up"


help_msg = """I answer to the following commands:

/start - I say hi.
/about - What I'm about.
/explain - I explain how I work.
/help - I send this message.
/count - I tell you how many messages from this chat I have processed.
/speak - Forces me to speak (if you are allowed to).
/period - Change the period of my messages. (Maximum of 100000)
/answer - Change my chance to answer at will. (Decimal between 0 and 1)
/restrict - Toggle restriction of configuration commands to admins only.
/silence - Toggle restriction on mentions by the bot.
/ping - Pong!
/where - Tell my configuration for this chat.
"""


about_msg = (
    "I am yet another Markov Bot experiment. I read everything you type to me"
    " and then spit back nonsensical messages that look like yours.\n\nYou can"
    " send /explain if you want further explanation."
)


explanation = (
    "I decompose every message I read in groups of 3 consecutive words, so for"
    " each consecutive pair I save the word that can follow them. I then use"
    " this to make my own messages. At first I will only repeat your messages"
    " because for each 2 words I will have very few possible following"
    " words.\n\nI also separate my vocabulary by chats, so anything I learn in"
    " one chat I will only say in that chat. For privacy, you know. Also, I"
    " save my vocabulary in the form of a json dictionary, so no logs are"
    " kept.\n\nMy default period in private chats is one message of mine from"
    " each 2 messages received, and in group chats it's 10 messages I read for"
    " each message I send."
)


def static_reply(text: str, format=None):
    async def reply(update: Update, context: ContextTypes.DEFAULT_TYPE):
        assert update.message
        await update.message.reply_text(text, parse_mode=format)

    return reply


async def stop(update: Update, context: ContextTypes.DEFAULT_TYPE):
    assert update.message
    chat = update.message.chat
    chat_id = str(chat.id)
    # TODO: remove chat from DB
    logger.warning(
        "I got blocked by {} [{}]".format(chat.effective_name, chat_id)
    )


# Legacy function from original velascobot, created before "effective_name"
# was implemented in the telegram library. Kept to ensure consistency
def get_chat_title(chat: Chat):
    """Returns group chat name or falls back to user name on a private chat."""
    if chat.title is not None:
        return chat.title
    elif chat.first_name is not None:
        if chat.last_name is not None:
            return chat.first_name + " " + chat.last_name
        else:
            return chat.first_name
    else:
        return ""


async def ping(update: Update, _context: ContextTypes.DEFAULT_TYPE) -> None:
    """Pongs back."""

    assert update.message is not None
    await update.message.reply_text("Pong!")


async def error_handler(update: object, context: ContextTypes.DEFAULT_TYPE):
    logger.error(context.error)
    sys.exit(1)


class Bot:
    def __init__(
        self,
        messenger: Messenger,
        bot_token: str,
        motd_chats: list[str] = [],
        motd_list: list[str] = [],
        whitelist: list[str] | None = None,
        testing: bool = False,
    ) -> None:
        self.messenger = messenger
        self.motd_chats = motd_chats
        self.motd_list = motd_list
        self.whitelist = whitelist
        self.testing = testing
        self.countdown: dict[str, int] = {}

        if testing:
            assert (
                whitelist is not None
            ), "Testing mode enabled but no whitelist was given."

        # Create the Application and pass it your bot's token.
        app = (
            Application.builder()
            .token(bot_token)
            .concurrent_updates(True)
            .build()
        )

        # on different commands - answer in Telegram
        app.add_handler(CommandHandler("start", static_reply(start_msg)))
        app.add_handler(CommandHandler("about", static_reply(about_msg)))
        app.add_handler(CommandHandler("explain", static_reply(explanation)))
        app.add_handler(CommandHandler("help", static_reply(help_msg)))
        app.add_handler(CommandHandler("stop", stop))
        app.add_handler(CommandHandler("ping", ping))
        app.add_handler(CommandHandler("count", self.get_processed_messages))
        app.add_handler(CommandHandler("speak", self.speak_command))
        app.add_handler(CommandHandler("period", self.period))
        app.add_handler(CommandHandler("answer", self.answer))
        app.add_handler(CommandHandler("restrict", self.restrict))
        app.add_handler(CommandHandler("silence", self.silence))
        app.add_handler(
            CommandHandler("where", static_reply("Not implemented"))
        )
        app.add_handler(
            MessageHandler(
                (
                    filters.TEXT
                    | filters.Sticker.ALL
                    | filters.ANIMATION
                    | filters.VIDEO
                ),
                self.read,
            )
        )

        app.add_error_handler(error_handler)

        self.app = app

    async def get_processed_messages(
        self, update: Update, _context: ContextTypes.DEFAULT_TYPE
    ) -> None:
        assert update.message

        chat_id = str(update.message.chat.id)
        chat_data = self.messenger.get_chat_data(chat_id)
        if chat_data is None:
            await update.message.reply_text(
                "I've processed no messages so far."
            )
            return

        num = (
            str(chat_data.processed_messages)
            if chat_data.processed_messages > 0
            else "no"
        )
        await update.message.reply_text(
            "I've processed {} messages so far.".format(num)
        )

    async def period(
        self, update: Update, _context: ContextTypes.DEFAULT_TYPE
    ) -> None:
        """Responds with the current message period or sets a new one."""
        assert update.message
        chat = update.message.chat

        chat_id = str(chat.id)
        chat_data = self.messenger.get_chat_data(chat_id)
        assert chat_data

        assert update.message.text
        words = update.message.text.split()
        if len(words) <= 1:
            await update.message.reply_text(
                f"The current speech period is {chat_data.message_period}"
            )
            return

        if chat_data.restricted:
            assert update.message.from_user
            user = await chat.get_member(update.message.from_user.id)
            if not self.superuser(user):
                await update.message.reply_text(
                    "You do not have permissions to do that."
                )
                return
        try:
            period = int(words[1])
            self.messenger.update_message_period(chat_id, period)
            await update.message.reply_text(
                f"Period of speaking set to {period}."
            )
        except Exception:
            await update.message.reply_text(
                "Format was confusing; period unchanged from"
                f" {chat_data.message_period}."
            )

    async def answer(
        self, update: Update, _context: ContextTypes.DEFAULT_TYPE
    ) -> None:
        """Responds with the current answer chance or sets a new one."""
        assert update.message
        chat = update.message.chat

        chat_id = str(chat.id)
        chat_data = self.messenger.get_chat_data(chat_id)
        assert chat_data

        assert update.message.text
        words = update.message.text.split()
        if len(words) <= 1:
            await update.message.reply_text(
                f"The current answer chance is {chat_data.answer_chance}"
            )
            return

        if chat_data.restricted:
            assert update.message.from_user
            user = await chat.get_member(update.message.from_user.id)
            if not self.superuser(user):
                await update.message.reply_text(
                    "You do not have permissions to do that."
                )
                return
        try:
            chance = float(words[1])
            self.messenger.update_answer_chance(chat_id, chance)
            await update.message.reply_text(
                f"Chance of answering set to {chance}."
            )
        except Exception:
            await update.message.reply_text(
                "Format was confusing; answer chance unchanged from"
                f" {chat_data.answer_chance}."
            )

    async def restrict(
        self, update: Update, _context: ContextTypes.DEFAULT_TYPE