openKriemy/openkriemy.py

from random import random, uniform
from asyncio import run, sleep
from datetime import datetime, timedelta, date, time

from telethon import TelegramClient
from telethon.events import NewMessage
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,
)
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()


async def markov_say(bot, peer_id, reply_to=None):
    if not markov.is_ready:
        return

    text = markov.generate()

    async with bot.action(peer_id, "typing"):
        amount = 0
        for _ in range(len(text)):
            amount += round(uniform(0.05, 0.2), 2)

        await sleep(min(amount, 8))

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


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


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

    try:
        command = parse_command(event.text)
    except ValueError:
        if await is_markov_enabled(peer_id):
            markov.extend_corpus(event.text)

            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")
            ) 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(random(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()