openKriemy/openkriemy.py

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 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
from markov import Markov

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


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


@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=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 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")

            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

    if not await is_allowed(peer_id):
        if 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

    try:
        action = await find_action(command.name)
    except SyntaxError:
        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:
        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:
                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)


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