diff --git a/bot/firestore.py b/bot/firestore.py new file mode 100644 index 0000000..fd540a0 --- /dev/null +++ b/bot/firestore.py @@ -0,0 +1,40 @@ +from bot.firestore_config import db + + +# New user setup +def new_user_setup(telegram_id, sheet_id): + user_ref = db.collection("users").document(str(telegram_id)) + user_ref.set({"sheet_id": sheet_id}) + + +# Check if user exists +def check_if_user_exists(telegram_id): + user_ref = db.collection("users").document(str(telegram_id)) + user_doc = user_ref.get() + return user_doc.exists + + +# Get user sheet id +def get_user_sheet_id(telegram_id): + user_ref = db.collection("users").document(str(telegram_id)) + user_doc = user_ref.get() + if user_doc.exists: + return user_doc.get("sheet_id") + else: + return None + + +# Get all user IDs +def get_all_user_id(): + users_ref = db.collection("users") + user_ids = [int(user.id) for user in users_ref.stream()] + return user_ids + + +# Get all sheet IDs +def get_all_sheet_id(): + users_ref = db.collection("users") + sheet_ids = [] + for user in users_ref.stream(): + sheet_ids.append(user.get("sheet_id")) + return sheet_ids diff --git a/bot/firestore_config.py b/bot/firestore_config.py new file mode 100644 index 0000000..4ad0fe5 --- /dev/null +++ b/bot/firestore_config.py @@ -0,0 +1,12 @@ +from firebase_admin import firestore +import firebase_admin +from firebase_admin import credentials +import os +import json + +firebase_json = json.loads(os.environ["FIREBASE_JSON"]) + +cred = credentials.Certificate(firebase_json) +firebase_admin.initialize_app(cred) + +db = firestore.client() diff --git a/bot/telegram_bot.py b/bot/telegram_bot.py index e7eca28..8931443 100644 --- a/bot/telegram_bot.py +++ b/bot/telegram_bot.py @@ -16,7 +16,7 @@ from bot.common import EntryType from bot.common import ConversationState as CS import bot.google_sheet as gs -import bot.firebase as db +import bot.firestore as db timezone = pytz.timezone("Asia/Singapore") @@ -37,6 +37,7 @@ def check_date_format(date_string): pattern = r"\b\d{1,2}\s(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\b" return bool(re.fullmatch(pattern, date_string, re.IGNORECASE)) + def check_month_format(month_string): pattern = r"^(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)$" return bool(re.fullmatch(pattern, month_string, re.IGNORECASE)) @@ -623,6 +624,7 @@ def get_day_transaction(update, context): update.message.reply_text(ERROR_TEXT) return ConversationHandler.END + def get_overall(update, context): context.user_data.clear() telegram_id = update.effective_user.id @@ -634,6 +636,7 @@ def get_overall(update, context): update.message.reply_text(ERROR_TEXT) return ConversationHandler.END + def handle_get_transaction(update, context): sheet_id = context.user_data["sheet_id"] reply = update.message.text @@ -662,21 +665,34 @@ def handle_get_transaction(update, context): except Exception as e: update.message.reply_text(ERROR_TEXT) + def handle_get_overall(update, context): sheet_id = context.user_data["sheet_id"] month = update.message.text - msg = "" try: if check_month_format(month): values = gs.get_overall(sheet_id, month) - msg = "```\n" - for row in values: - if len(row) == 3: - msg += "{:<20} {:<10} {:<10}\n".format(*row) - elif len(row) == 2: - msg += "{:<20} {:<10}\n".format(*row) - msg += "```" - update.message.reply_text(msg, parse_mode='Markdown') + + final_income = values[0] + msg = f"--- Final Income ---\n`{final_income[2]}`\n\n" + + msg += "--- Spending ---\n" + max_len = ( + max(len(row[0]) for row in values[1:-2]) + 1 + ) # +1 for some extra space + for row in values[1:-2]: + if row[1].startswith("-"): # if value is negative + msg += f"`{row[0].ljust(max_len)}{row[1]}`\n" + else: + msg += f"`{row[0].ljust(max_len)} {row[1]}`\n" # else keep the original space + + total_spent = values[-2] + msg += f"\n--- Total Spent ---\n`{total_spent[1]}`\n" + + overall = values[-1] + msg += f"\n--- Overall ---\n`{overall[2]}`\n" + + update.message.reply_text(msg, parse_mode="Markdown") return ConversationHandler.END else: update.message.reply_text(GET_OVERALL_TEXT) @@ -684,6 +700,7 @@ def handle_get_overall(update, context): except Exception as e: update.message.reply_text(ERROR_TEXT) + def add_income(update, context): context.user_data.clear() telegram_id = update.effective_user.id @@ -807,7 +824,7 @@ def setup_handlers(dispatcher): CommandHandler("addothers", add_others), CommandHandler("addincome", add_income), CommandHandler("getdaytransaction", get_day_transaction), - CommandHandler("getoverall", get_overall) + CommandHandler("getoverall", get_overall), ], states={ CS.SET_UP: [MessageHandler(Filters.text & ~Filters.command, set_up)],