From c0b7edede62b9fc9cb7aa9364e1c5b871596ff6a Mon Sep 17 00:00:00 2001
From: juckerj <jonas.jucker@env.ethz.ch>
Date: Sun, 19 Feb 2023 15:28:11 +0100
Subject: [PATCH] draft new subscription mode

---
 bot.py | 107 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 106 insertions(+), 1 deletion(-)

diff --git a/bot.py b/bot.py
index a4a3889..17a0dc4 100644
--- a/bot.py
+++ b/bot.py
@@ -2,15 +2,19 @@
 import socket
 import time
 
-from telegram import ReplyKeyboardMarkup, Update
+from telegram import ReplyKeyboardMarkup, Update, ReplyKeyboardRemove
 from telegram.ext import (
     CommandHandler,
+    MessageHandler,
+    Filters,
     ConversationHandler,
     Updater,
     PicklePersistence,
     CallbackContext,
 )
 
+CHOOSE, STATION = range(2)
+
 class PlotBot:
 
     def __init__(self,token):
@@ -24,12 +28,104 @@ def __init__(self,token):
         self._dp.add_handler(CommandHandler('subscribe',self._subscribe))
         self._dp.add_handler(CommandHandler('unsubscribe',self._unsubscribe))
         self._dp.add_handler(CommandHandler('where_am_I',self._get_ip_address))
+        self._dp.add_handler(CommandHandler('conv',self._conv))
+
+
+        subscription_handler = ConversationHandler(
+            entry_points=[MessageHandler(Filters.regex('^(subscribe)$'), self._choose_station)],
+            states={
+                STATION: [MessageHandler(Filters.regex('^(Davos|Tschiertschen)$'), self._subscribe_for_station)],
+                },
+            fallbacks=[CommandHandler('cancel', self._cancel)],
+            )
+
+        unsubscription_handler = ConversationHandler(
+            entry_points=[MessageHandler(Filters.regex('^(unsubscribe)$'), self._unchoose_station)],
+            states={
+                STATION: [MessageHandler(Filters.regex('^(Davos|Tschiertschen)$'), self._unsubscribe_for_station)],
+                },
+            fallbacks=[CommandHandler('cancel', self._cancel)],
+            )
+
+        self._dp.add_handler(subscription_handler)
+        self._dp.add_handler(unsubscription_handler)
 
         # start the bot
         self.updater.start_polling()
 
         self._new_users_waiting_for_plots = []
 
+
+    def _conv(self,update: Update, context: CallbackContext):
+        reply_keyboard = [['subscribe', 'unsubscribe']]
+
+        reply_text = "Hi! I am OpenEns. I supply you with the latest ECWMF meteograms. \
+                      As soon as the latest forecast is available I deliver them to you. \
+                      Choose one of the following actions below."
+        update.message.reply_text(reply_text,
+            reply_markup=ReplyKeyboardMarkup(reply_keyboard, one_time_keyboard=True),
+        )
+
+    def _unchoose_station(self,update: Update, context: CallbackContext) -> int:
+        reply_keyboard = [['Davos', 'Tschiertschen']]
+
+        reply_text = "Choose a station"
+        update.message.reply_text(reply_text,
+            reply_markup=ReplyKeyboardMarkup(reply_keyboard, one_time_keyboard=True),
+        )
+
+        return STATION
+
+    def _choose_station(self,update: Update, context: CallbackContext) -> int:
+        reply_keyboard = [['Davos', 'Tschiertschen']]
+
+        reply_text = "Choose a station"
+        update.message.reply_text(reply_text,
+            reply_markup=ReplyKeyboardMarkup(reply_keyboard, one_time_keyboard=True),
+        )
+
+        return STATION
+
+    def _unsubscribe_for_station(self,update: Update, context: CallbackContext) -> int:
+        user = update.message.from_user
+        msg_text = update.message.text
+        reply_text = f'Unubscribed for Station {msg_text}'
+        update.message.reply_text(reply_text,
+        reply_markup=ReplyKeyboardRemove(),
+        )
+        self._revoke_subscription(user,msg_text,context)
+        logging.info(f' {user.first_name} unsubscribed for Station {msg_text}')
+
+        return ConversationHandler.END
+
+    def _subscribe_for_station(self,update: Update, context: CallbackContext) -> int:
+        user = update.message.from_user
+        msg_text = update.message.text
+        reply_text = f'Subscribed for Station {msg_text}'
+        update.message.reply_text(reply_text,
+        reply_markup=ReplyKeyboardRemove(),
+        )
+        self._register_subscription(user,msg_text,context)
+        logging.info(f' {user.first_name} subscribed for Station {msg_text}')
+
+        return ConversationHandler.END
+
+    def _register_subscription(self,user,station,context):
+
+        # add user to subscription list for given station
+        context.bot_data.setdefault(station, set())
+        context.bot_data[station].add(user.id)
+
+        logging.info(context.bot_data.setdefault(station, set()))
+
+    def _revoke_subscription(self,user,station,context):
+
+        # remove user to subscription list for given station
+        if user.id in context.bot_data[station]:
+            context.bot_data[station].remove(user.id)
+
+        logging.info(context.bot_data.setdefault(station, set()))
+
     def _start(self,update: Update, context: CallbackContext):
         reply_text = "Hi! I am OpenEns. I supply you with the latest ECWMF meteograms. \
                       As soon as the latest forecast is available I deliver them to you. \
@@ -79,6 +175,15 @@ def _get_ip_address(self,update: Update, context: CallbackContext):
         reply_text = f"IP-ADDRESS: {ip_address}"
         update.message.reply_text(reply_text)
 
+    def _cancel(self,update: Update, context: CallbackContext) -> int:
+        user = update.message.from_user
+        logging.info("User %s canceled the conversation.", user.first_name)
+        update.message.reply_text(
+            'Bye! I hope we can talk again some day.', reply_markup=ReplyKeyboardRemove()
+        )
+
+        return ConversationHandler.END
+
     def has_new_subscribers_waiting(self):
         if self._new_users_waiting_for_plots:
             return True