From d246294c5fc72a214a4f9609ff198e66f05eee30 Mon Sep 17 00:00:00 2001
From: Steven Haussmann <hausss@rpi.edu>
Date: Wed, 1 Jul 2020 11:30:38 -0400
Subject: [PATCH] Get polling working

---
 icon-poller.py | 110 ++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 94 insertions(+), 16 deletions(-)

diff --git a/icon-poller.py b/icon-poller.py
index 06230a9..fd27539 100644
--- a/icon-poller.py
+++ b/icon-poller.py
@@ -9,6 +9,8 @@ import random
 
 submissions = {}
 names = {}
+polls = {}
+poll_options = {}
 status = {}
 
 def ensure_chat(func):
@@ -17,7 +19,9 @@ def ensure_chat(func):
         if chat_id not in submissions:
             submissions[chat_id] = {}
             names[chat_id] = {}
-            status[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
@@ -42,11 +46,31 @@ def update_name(func):
         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, *args, **kwargs):
+    def wrapper(update, context, *args, **kwargs):
         chat_id = update.effective_chat.id
         if status[chat_id] == "open":
-            return func(update, *args, **kwargs)
+            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
@@ -61,24 +85,76 @@ def accept_photo(update: Update, context: CallbackContext):
     print(file_id)
 
 @ensure_chat
-@ensure_user
-def echo(update: Update, context: CallbackContext):
-    context.bot.send_message(chat_id=update.effective_chat.id, text="Hi.")
-    updater.stop()
-
-@ensure_chat
+@ensure_closed
 def open_submissions(update: Update, context: CallbackContext):
     chat_id = update.effective_chat.id
-    context.bot.send_message(chat_id=update.effective_chat.id, text="Send me your images pls.")
+    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_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.choices(users, k=min(len(users), 10))
+
+    options = []
+    chosen = []
+
+    for user_id in chosen_users:
+        options.append(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"
+    if 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 file_id: InputMediaPhoto(file_id), 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
-def show_submissions(update: Update, context: CallbackContext):
+@ensure_polling
+def decide_poll(update: Update, context: CallbackContext):
     chat_id = update.effective_chat.id
-    options = list(map(lambda file_id: InputMediaPhoto(file_id), submissions[chat_id].values()))
-    chosen = random.choices(options, k=min(len(options), 10))
-    context.bot.send_media_group(chat_id=chat_id, media=[chosen[0]])
-    
+    status[chat_id] = "closed"
+
+    result = context.bot.stop_poll(chat_id=chat_id, message_id=polls[chat_id])
+
+    votes = -1
+    winner = 0
+
+    for index, option in enumerate(result.options):
+        if option.voter_count > votes:
+            winner = index
+            votes = option.voter_count
+
+    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
+def cancel(update: Update, context: CallbackContext):
+    chat_id = update.effective_chat.id
+    status[chat_id] = "closed"
+    submissions[chat_id] = {}
+    del polls[chat_id]
+
 
 if __name__ == "__main__":
     try:
@@ -92,7 +168,9 @@ if __name__ == "__main__":
 
     dispatcher.add_handler(MessageHandler(Filters.photo, accept_photo))
     dispatcher.add_handler(CommandHandler("open", open_submissions))
-    dispatcher.add_handler(CommandHandler("show", show_submissions))
+    dispatcher.add_handler(CommandHandler("poll", start_poll))
+    dispatcher.add_handler(CommandHandler("decide", decide_poll))
+    dispatcher.add_handler(CommandHandler("cancel", cancel))
 
     updater.start_polling()
     updater.idle()
\ No newline at end of file