initial fcm cloud messaging trigger in generate script

5 files changed, 117 insertions, 2 deletions

diff --git a/.env.template b/.env.template
index d39d675..bef2005 100644
--- a/.env.template
+++ b/.env.template
@@ -1,3 +1,8 @@
 DISCORD_AUTHORIZATION=
 OPENAI_API_KEY=
 GOOGLE_TRANSLATE_API_KEY=
+
+# Optional for FCM Notifications
+KV_REST_API_TOKEN=
+KV_REST_API_URL=
+FIREBASE_SERVICE_ACCOUNT_JSON_PATH=serviceAccount.json
diff --git a/.gitignore b/.gitignore
index 9037223..2dc61e1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -176,3 +176,4 @@ wac_result_cache.json
 summarization_cache.json
 *.bak
 *.db
+*.json
diff --git a/database.py b/database.py
index 2f6ae9a..f23848d 100644
--- a/database.py
+++ b/database.py
@@ -147,3 +147,14 @@ class Database:
             return None
         is_news = True if result[0] == 1 else False
         return is_news, result[1]
+
+    def check_news_id_exists(self, key: str) -> bool:
+        """
+        Check if a news entry with the given ID exists in the database.
+
+        :param key: The ID of the news entry to check.
+        :return: True if the news entry exists, False otherwise.
+        """
+        self._cursor.execute("SELECT 1 FROM news WHERE news_id = ?", (key,))
+        result = self._cursor.fetchone()
+        return result is not None
diff --git a/generate.py b/generate.py
index f5c6b37..5ffa464 100644
--- a/generate.py
+++ b/generate.py
@@ -12,12 +12,14 @@ from dotenv import load_dotenv
 from datetime import datetime, timedelta
 from database import Database
 from feed import build_rss_from_news_feed
+from notifications import check_can_send_notifs, broadcast_to_topic
 
 load_dotenv()
 
 OUTPUT_DIR = "news"
 ARCHIVE_NEWS = True
 GENERATE_RSS_FEEDS = True
+SEND_NOTIFICATIONS = True
 
 
 def compute_json_hash(data):
@@ -61,6 +63,35 @@ def attach_news_meta_data(news_data: list):
         "news_posts": news_data,
     }
 
+def attempt_broadcast_notifications(news_data: list, title: str, topic: str, image: str="http://arcade.moekyun.me/rasis.webp", limit=constants.DAYS_LIMIT):
+    if news_data and SEND_NOTIFICATIONS:
+        if not check_can_send_notifs:
+            print("[WARNING] Skipping notifications as env vars are not properly configured. See template")
+        else:
+            database = Database()
+            cutoff = datetime.now() - timedelta(days=limit)
+            for entry in news_data:
+                if datetime.fromtimestamp(entry["timestamp"]) < cutoff:
+                    continue
+                news_id = compute_json_hash(entry)
+                if database.check_news_id_exists(news_id):
+                    continue
+                else:
+                    print("[Notifications] Broadcasting Unique News with id: " + news_id)
+                    input()
+                    if entry.get("en_headline"):
+                        body_text = entry["en_headline"]
+                    elif entry.get("en_content"):
+                        body_text = entry["en_content"]
+                    elif entry.get("headline"):
+                        body_text = entry["headline"]
+                    else:
+                        body_text = entry.get("content", "")
+                    if len(body_text) > 50:
+                        body_text = body_text[:50] + "..."
+                    broadcast_to_topic(topic, title, body_text, image)
+
+
 
 def log_output(message: str, type: str = "DEBUG"):
     """
@@ -119,7 +150,9 @@ def generate_iidx_news_file(eamuse_feed: bool=False):
         return generate_news_file("iidx_news", constants.IIDX_PINKY_CRUSH_NEWS_SITE)
 
 def generate_sdvx_news_file():
-    return generate_news_file("sdvx_news", constants.SOUND_VOLTEX_EXCEED_GEAR_NEWS_SITE)
+    news = generate_news_file("sdvx_news", constants.SOUND_VOLTEX_EXCEED_GEAR_NEWS_SITE)
+    attempt_broadcast_notifications(news, "New Information for SOUND VOLTEX","sdvx")
+    return news
 
 def generate_ddr_news_file(eamuse_feed: bool=False):
     if eamuse_feed:
@@ -195,9 +228,9 @@ if __name__ == "__main__":
     if not os.path.exists(OUTPUT_DIR):
         log_output(f"{OUTPUT_DIR} was not found. Creating this directory...")
         os.makedirs(OUTPUT_DIR)
+    sdvx_news_data = generate_sdvx_news_file()
     polaris_news_data = generate_polaris_chord_news_file()
     iidx_news_data = generate_iidx_news_file(eamuse_feed=True)
-    sdvx_news_data = generate_sdvx_news_file()
     ddr_news_data = generate_ddr_news_file(eamuse_feed=True)
     dance_rush_news_data = generate_dance_rush_news_file()
     dance_around_news_data = generate_dance_around_news_file()
diff --git a/notifications.py b/notifications.py
new file mode 100644
index 0000000..f5c7e77
--- /dev/null
+++ b/notifications.py
@@ -0,0 +1,65 @@
+import firebase_admin
+from firebase_admin import credentials, messaging
+from dotenv import load_dotenv
+import requests
+import os
+
+load_dotenv()
+
+def check_can_send_notifs():
+    required_env_vars = [
+        "KV_REST_API_URL",
+        "KV_REST_API_TOKEN",
+        "FIREBASE_SERVICE_ACCOUNT_JSON_PATH"
+    ]
+    missing_vars = [var for var in required_env_vars if os.environ.get(var) is None]
+    if missing_vars:
+        return False
+    else:
+        return True
+
+def get_tokens(topic):
+    upstash_rest_api_url = os.environ.get("KV_REST_API_URL")
+    upstash_token = os.environ.get("KV_REST_API_TOKEN")
+    url = f"{upstash_rest_api_url}/smembers/fcm-{topic}"
+    resp = requests.get(
+        url,
+        headers={"Authorization": f"Bearer {upstash_token}"}
+    )
+    if resp.status_code != 200:
+        print("Upstash error:", resp.text)
+        return []
+    data = resp.json()
+    return data.get("result", [])
+
+def broadcast_to_topic(topic: str, title: str, body: str, image):
+    if not firebase_admin._apps:
+        cred = credentials.Certificate(os.environ.get("FIREBASE_SERVICE_ACCOUNT_JSON_PATH"))
+        firebase_admin.initialize_app(cred)
+
+    tokens = get_tokens(topic)
+    if not tokens:
+        print(f"No tokens found for topic '{topic}'")
+        return
+
+    for token in tokens:
+        message = messaging.Message(
+            notification=messaging.Notification(
+                title=title,
+                body=body,
+                image=image
+            ),
+            token=token
+        )
+        try:
+            response = messaging.send(message)
+            print(f"Sent to {token}: {response}")
+        except Exception as e:
+            print(f"Error sending to {token}: {e}")
+
+if __name__ == "__main__":
+    topic = "sdvx"
+    title = "New sdvx Update!"
+    body = "Check out the latest DDR content now!"
+
+    broadcast_to_topic(topic, title, body)