icon-poller/icon-poller.py

from telegram.ext import Updater, MessageHandler, Filters, CallbackContext, CommandHandler
from telegram import Update, InputMediaPhoto
import json
import sys
import io
import logging
from PIL import Image
import random

from mwt import MWT

submissions = {}
names = {}
polls = {}
poll_options = {}
status = {}

@MWT(timeout=60*60)
def get_admin_ids(bot, chat_id):
    """Returns a list of admin IDs for a given chat. Results are cached for 1 hour."""
    return [admin.user.id for admin in bot.get_chat_administrators(chat_id)]

def ensure_admin(func):
    def wrapper(update, context, *args, **kwargs):
        chat_id = update.effective_chat.id
        user_id = update.effective_user.id
        if user_id in get_admin_ids(context.bot, chat_id):
            return func(update, context, *args, **kwargs)
    return wrapper

def ensure_chat(func):
    def wrapper(update, *args, **kwargs):
        chat_id = update.effective_chat.id
        if chat_id not in submissions:
            submissions[chat_id] = {}
            names[chat_id] = {}
            status[chat_id] = "closed"
            polls[chat_id] = {}
            poll_options[chat_id] = []
            print("Added chat: ", chat_id)
        return func(update, *args, **kwargs)
    return wrapper

def ensure_user(func):
    def wrapper(update, *args, **kwargs):
        chat_id = update.effective_chat.id
        user_id = update.effective_user.id
        if user_id not in submissions[chat_id]:
            names[chat_id][user_id] = update.effective_user.first_name
            print("Added user: ", user_id, " - ", update.effective_user.first_name)
        return func(update, *args, **kwargs)
    return wrapper

def update_name(func):
    def wrapper(update, *args, **kwargs):
        chat_id = update.effective_chat.id
        user_id = update.effective_user.id
        if user_id not in submissions[update.effective_chat.id]:
            names[chat_id][user_id] = update.effective_user.first_name
            print("Got name for ", user_id, ": ", update.effective_user.first_name)
        return func(update, *args, **kwargs)
    return wrapper

def ensure_closed(func):
    def wrapper(update, context, *args, **kwargs):
        chat_id = update.effective_chat.id
        if status[chat_id] == "closed":
            return func(update, context, *args, **kwargs)
        else:
            context.bot.send_message(chat_id=chat_id, text="I need to be closed to do that. I'm currently: " + status[chat_id])
    return wrapper

def ensure_open(func):
    def wrapper(update, context, *args, **kwargs):
        chat_id = update.effective_chat.id
        if status[chat_id] == "open":
            return func(update, context, *args, **kwargs)
        else:
            context.bot.send_message(chat_id=chat_id, text="I need to be open to do that. I'm currently: " + status[chat_id])
    return wrapper

def ensure_polling(func):
    def wrapper(update, context, *args, **kwargs):
        chat_id = update.effective_chat.id
        if status[chat_id] == "polling":
            return func(update, context, *args, **kwargs)
        else:
            context.bot.send_message(chat_id=chat_id, text="I need to be polling to do that. I'm currently: " + status[chat_id])
    return wrapper

@ensure_chat
@ensure_user
@update_name
@ensure_open
def accept_photo(update: Update, context: CallbackContext):
    chat_id = update.effective_chat.id
    user_id = update.effective_user.id
    file_id = update.message.photo[-1].file_id
    submissions[chat_id][user_id] = file_id
    print(file_id)

@ensure_chat
@ensure_admin
@ensure_closed
def open_submissions(update: Update, context: CallbackContext):
    chat_id = update.effective_chat.id
    submissions[chat_id] = {}
    context.bot.send_message(chat_id=chat_id, text="Send your images as a reply to this post.")
    status[chat_id] = "open"
    
@ensure_chat
@ensure_admin
@ensure_open
def start_poll(update: Update, context: CallbackContext):
    chat_id = update.effective_chat.id
    status[chat_id] = "polling"
    poll_options[chat_id] = []

    users = list(submissions[chat_id].keys())

    print("Users with submissions: ", users)

    chosen_users = random.sample(users, k=min(len(users), 10))

    options = []
    chosen = []

    for index, user_id in enumerate(chosen_users):
        options.append(str(index+1) + ": " + names[chat_id][user_id])
        chosen.append(submissions[chat_id][user_id])
        poll_options[chat_id].append(user_id)

    if len(chosen) == 0:
        context.bot.send_message(chat_id=chat_id, text="There were no submissions.")
        status[chat_id] = "closed"
    elif len(chosen) == 1:
        context.bot.send_photo(chat_id=chat_id, photo=chosen[0])
        context.bot.send_message(chat_id=chat_id, text="There was only one submission. " + options[0] + " is the winner.")
        status[chat_id] = "closed"
    else:
        input_media = list(map(lambda pair: InputMediaPhoto(pair[1], caption=pair[0]), zip(options, chosen)))
        context.bot.send_media_group(chat_id=chat_id, media=input_media)    

        polls[chat_id] = context.bot.send_poll(chat_id=chat_id, options = options, question="Pick an icon").message_id

@ensure_chat
@ensure_admin
@ensure_polling
def decide_poll(update: Update, context: CallbackContext):
    chat_id = update.effective_chat.id
    status[chat_id] = "closed"

    result = context.bot.stop_poll(chat_id=chat_id, message_id=polls[chat_id])

    votes = -1
    winners = []

    for index, option in enumerate(result.options):
        if option.voter_count > votes:
            winners = []
        if option.voter_count >= votes:
            winners.append(index)
            votes = option.voter_count

    winner = random.choice(winners)
    winner_id = poll_options[chat_id][winner]
    winner_name = names[chat_id][winner_id]
    winner_photo = submissions[chat_id][winner_id]

    context.bot.send_message(chat_id=chat_id, text="The winner was " + winner_name)
    context.bot.send_photo(chat_id=chat_id, photo=winner_photo)

@ensure_chat
@ensure_admin
def cancel(update: Update, context: CallbackContext):
    chat_id = update.effective_chat.id
    status[chat_id] = "closed"
    submissions[chat_id] = {}
    del polls[chat_id]

@ensure_chat
def show_status(update: Update, context: CallbackContext):
    chat_id = update.effective_chat.id
    if status[chat_id] == "closed":
        context.bot.send_message(chat_id=chat_id, text="I am currently closed.")
    elif status[chat_id] == "open":
        context.bot.send_message(chat_id=chat_id, text=f"I am currently open for submissions. There are {len(submissions[chat_id])} submissions.")
    elif status[chat_id] == "polling":
        context.bot.send_message(chat_id=chat_id, text="I am currently polling.")


if __name__ == "__main__":
    try:
        config = json.load(open("config.json", "r", encoding="utf-8"))
    except:
        logging.error("Couldn't read the config!")
        sys.exit(1)

    updater = Updater(token=config["token"], use_context=True)
    dispatcher = updater.dispatcher

    dispatcher.add_handler(MessageHandler(Filters.photo, accept_photo))
    dispatcher.add_handler(CommandHandler("open", open_submissions))
    dispatcher.add_handler(CommandHandler("poll", start_poll))
    dispatcher.add_handler(CommandHandler("decide", decide_poll))
    dispatcher.add_handler(CommandHandler("cancel", cancel))
    dispatcher.add_handler(CommandHandler("status", show_status))


    updater.start_polling()
    updater.idle()