openKriemy/openkriemy.py

import re
from random import random, randint
from asyncio import sleep
from datetime import datetime, timedelta, date, time

from telethon import TelegramClient
from telethon.events import NewMessage, InlineQuery
from telethon.utils import resolve_bot_file_id, get_peer_id

from tortoise.exceptions import ConfigurationError

from actions import get_all_birthdays
from utils import parse_command, get_link_to_user, calculate_age, Kind
from config import config
from db import init_db
from actions import (
    find_action,
    get_random_gif,
    is_admin,
    is_allowed,
    is_markov_enabled,
    get_markov_option,
    list_markov_chats,
    markov_say,
    run,
)
from commands import COMMANDS, CHESS_COMMANDS, CHESS_ALIASES
from markov import Markov
from chess0 import ChessManager

bot = TelegramClient("openkriemy", config.API_ID, config.API_HASH).start(
    bot_token=config.API_TOKEN
)
markov = Markov()

if config.CHESS_UCI_ENGINE:
    chess = ChessManager(config.CHESS_UCI_ENGINE)
else:
    chess = None


# Wait isn't that illegal??
bot.markov = markov
bot.chess = chess


@bot.on(InlineQuery)
async def on_inline_query(event):
    text = event.text.strip()
    if not text:
        return

    lang, input_text, result = await run(text)
    if not lang:
        await event.answer(
            [
                event.builder.article(result, text=result),
            ]
        )

        return

    await event.answer(
        [
            event.builder.article(
                "Результат.",
                text=f"Язык: {lang}\nКод:```\n{input_text.replace('`', '')}```Результат: {result}",
            ),
        ]
    )


@bot.on(NewMessage)
async def on_message(event):
    peer_id = get_peer_id(event.peer_id)
    text = event.text

    try:
        command = parse_command(text)
    except ValueError:
        if chess and event.sender and (event.sender.id in chess.sessions or str(peer_id) in chess.sessions):
            reply = await event.get_reply_message()
            if reply and get_peer_id(reply.from_id) == await bot.get_peer_id("me") and reply.media:
                text = text.strip()
                cmd = text.lower().split(" ")[0]

                if cmd not in CHESS_COMMANDS and cmd not in CHESS_ALIASES:
                    text = f"/chess m {text}"
                else:
                    text = f"/chess {text}"

                await COMMANDS["chess"].handler(bot, event, parse_command(text))

                return

        if await is_markov_enabled(peer_id):
            markov.extend_corpus(text)

            for word in config.MARKOV_TRIGGER_WORDS:
                if word.lower() in text.lower() and random() > 0.5:
                    await markov_say(bot, peer_id, reply_to=event)

                    return

            reply_prob = await get_markov_option(peer_id, "opt_reply_prob")
            if reply_prob < 0:
                return

            reply = await event.get_reply_message()
            if (
                reply
                and get_peer_id(reply.from_id) == await bot.get_peer_id("me")
                and random() > 0.5
            ) or random() > reply_prob:
                await markov_say(bot, peer_id, reply_to=event)

        return

    handler = COMMANDS.get(command.name, None)

    if handler and handler.is_public:
        await handler.handler(bot, event, command)

        return

    action = None

    try:
        action = await find_action(command.name)
    except (SyntaxError, ConfigurationError):
        pass

    if not await is_allowed(peer_id):
        if not action and (not handler or not handler.is_restricted):
            return

    if handler:
        if handler.is_restricted and not await is_admin(bot, event.sender):
            await event.reply("К сожалению, данная команда Вам недоступна.")
        else:
            await handler.handler(bot, event, command)

        return

    if not action:
        return

    reply_to = None
    target = None

    if action.kind != Kind.NO_TARGET:
        target = await event.get_reply_message()

        if not target:
            try:
                target = await bot.get_entity(command.args[0])
            except (ValueError, IndexError):
                if action.kind != Kind.NO_TARGET_MAYBE:
                    await event.reply("Это действие нужно применить на кого-то!")

                    return
        else:
            reply_to = target
            target = target.sender

        if target is None:
            target = await bot.get_entity(event.peer_id.channel_id)

        if action.kind == Kind.CANNOT_APPLY_TO_SELF and target.id == event.sender.id:
            await event.reply("Данное действие нельзя применять к самому себе...")

            return

    try:
        await event.delete()
    except Exception:
        pass

    if event.sender is None:
        initiator = await bot.get_entity(event.peer_id.channel_id)
        initiator = initiator.title
    else:
        initiator = get_link_to_user(event.sender)

    text = action.template.format(
        **{"initiator": initiator, "target": get_link_to_user(target) if target else ""}
    )

    gif = await get_random_gif(action)

    if gif:
        gif = resolve_bot_file_id(gif.file_id)

    await bot.send_message(event.peer_id, message=text, file=gif, reply_to=reply_to)


async def notify_birthdays():
    birthdays = await get_all_birthdays()

    for birthday in birthdays:
        age = calculate_age(birthday.date)

        if age.days_until < 1:
            try:
                try:
                    entity = await bot.get_entity(birthday.user_id)
                except ValueError:
                    await bot.get_participants(birthday.peer_id)

                    entity = await bot.get_entity(birthday.user_id)

                await bot.send_message(
                    birthday.peer_id,
                    f"{get_link_to_user(entity)}, поздравляю с днём рождения!!~~",
                )
            except Exception:
                pass


async def notify_birthdays_loop():
    interval = datetime.combine(date.today(), time(hour=0, minute=0))

    while True:
        await sleep(((interval - datetime.now()) % timedelta(days=1)).total_seconds())

        await notify_birthdays()


async def markov_say_loop():
    while True:
        await sleep(randint(30, 60 * 5))

        for chat in await list_markov_chats():
            if random() > chat.opt_message_prob:
                await markov_say(bot, chat.peer_id)


async def chess_cleanup():
    while True:
        await sleep(60 * 60)
        await chess.cleanup()


with bot:
    bot.loop.run_until_complete(init_db())
    bot.loop.create_task(notify_birthdays_loop())
    bot.loop.create_task(markov_say_loop())
    if chess:
        bot.loop.create_task(chess_cleanup())
    bot.start()
    bot.run_until_disconnected()