openKriemy/actions.py

from random import randint

from tortoise.contrib.postgres.functions import Random
from telethon.utils import get_input_document
from telethon.tl.functions.stickers import CreateStickerSetRequest, AddStickerToSetRequest
from telethon.tl.functions.messages import UploadMediaRequest, GetStickerSetRequest
from telethon.tl.types import InputStickerSetID, InputStickerSetShortName, InputStickerSetItem, InputMediaUploadedDocument, InputPeerSelf
from tortoise.expressions import F

from models import Action, Gif, UserColor, StickerPack, Admin
from utils import is_valid_name
from config import config

async def is_admin(bot, user):
  admin = await bot.get_entity(config.ADMIN)

  if user.id == admin.id:
    return True

  admin = await Admin.filter(user_id=user.id)
  if admin:
    return True

  return False

async def add_admin(user):
  await Admin(user_id=user.id).save()

async def delete_admin(user):
  admin = await Admin.filter(user_id=user.id).first()
  if not admin:
    raise IndexError

  await admin.delete()

async def create_action(name, template, can_apply_to_self):
  if not is_valid_name(name):
    raise SyntaxError

  await Action(name=name, template=template, can_apply_to_self=can_apply_to_self).save()
   
async def find_action(name):
  if not is_valid_name(name):
    raise SyntaxError

  return await Action.filter(name=name).first()

async def delete_action(name):
  action = await find_action(name)
  if not action:
    raise NameError

  gifs = await action.gifs.all()
  for gif in gifs:
    await gif.delete()

  await action.delete()

async def add_gif(action, file_id):
  await Gif(action=action, file_id=file_id).save()

async def get_random_gif(action):
  return await action.gifs.all().annotate(order=Random()).order_by('order').first()

async def assign_color(username):
  user_color = await UserColor.filter(username=username).first()
  if not user_color:
    user_color = UserColor(username=username, color=randint(0, 7))
    await user_color.save()
  
  return user_color.color

async def create_new_pack(bot, sticker):
  last_pack = await StickerPack.all().order_by('-id').first()
  set_id = last_pack.id + 1 if last_pack else 1
 
  user = await bot.get_entity(config.USER)
  
  pack = await bot(CreateStickerSetRequest(
    user_id=user.id,
    title=f'Messages #{set_id}.',
    short_name=f'messages{set_id}_by_openkriemybot',
    stickers=[sticker]
  ))

  sid = pack.set.id
  hash = pack.set.access_hash

  await StickerPack(
    short_name=pack.set.short_name,
    sid=sid,
    hash=hash
  ).save()

  return pack

async def get_current_pack(bot):
  pack = await StickerPack.all().order_by('-id').first()
  if not pack or pack.stickers_count == 120:
    return None

  return pack

async def add_sticker(bot, file, emoji):
  file = await bot.upload_file(file)
  file = InputMediaUploadedDocument(file, 'image/png', [])
  file = await bot(UploadMediaRequest(InputPeerSelf(), file))
  file = get_input_document(file)
  sticker = InputStickerSetItem(document=file, emoji=emoji)
  
  pack = await get_current_pack(bot)
  if not pack:
    pack = await create_new_pack(bot, sticker)
  else:
    await StickerPack.filter(id=pack.id).update(stickers_count=F('stickers_count') + 1)

    pack = await bot(AddStickerToSetRequest(
      stickerset=InputStickerSetID(
        id=pack.sid,
        access_hash=pack.hash
      ), 
      sticker=sticker
    ))
    
  return get_input_document(pack.documents[-1])