From 06ba8bec34aa693fc69c0a749f72da667f650d0f Mon Sep 17 00:00:00 2001
From: Bulat Kurbanov <kurbanovbul@gmail.com>
Date: Wed, 6 Nov 2024 19:42:31 +0100
Subject: [PATCH] Add roles

---
 configs/hafmc.toml           | 20 ++++++++++++++---
 src/config.py                |  1 +
 src/services/notification.py | 42 +++++++++++++++++++++++++++++++++++-
 src/services/twitch.py       | 14 ++----------
 4 files changed, 61 insertions(+), 16 deletions(-)

diff --git a/configs/hafmc.toml b/configs/hafmc.toml
index 7e6cc38..daf79d7 100644
--- a/configs/hafmc.toml
+++ b/configs/hafmc.toml
@@ -2,16 +2,19 @@
 id = 59900845
 name = "hafmc"
 
-
 [notifications]
 start_stream = '''
-HafMC сейчас стримит {title} ({category})!
+HafMC стримит {title} ({category})!
 Присоединяйся: https://twitch.tv/hafmc
+
+{role}
 '''
 
 change_category = '''
-HafMC начал играть в {category}!
+HafMC играет в {category} ({title})!
 Присоединяйся: https://twitch.tv/hafmc
+
+{role}
 '''
 
 
@@ -25,5 +28,16 @@ notifications_channel_id = 1198296540964475002
 channel_id = 1201810638800691210
 message_id = 1239664178038313012
 
+[integrations.discord.roles]
+"Counter-Strike" = 1198372462476398823
+"SIGame" = 1198372523046346833
+"Battlefield 1" = 1198377604533735454
+"Battlefield V" = 1198377604533735454
+"Among Us" = 1198377604533735454
+"BattleBit Remastered" = 1198381867343286352
+"Lethal Company" = 1198383849785270302
+"HELLDIVERS 2" = 1230292196805054488
+"Satisfactory" = 1292091826802528308
+
 [integrations.telegram]
 notifications_channel_id = -1001939021131
diff --git a/src/config.py b/src/config.py
index baf1ccd..f568dc8 100644
--- a/src/config.py
+++ b/src/config.py
@@ -21,6 +21,7 @@ class DiscordConfig(BaseModel):
     guild_id: int
     notifications_channel_id: int
     games_list: GamesListConfig | None = None
+    roles: dict[str, int] | None = None
 
 class TelegramConfig(BaseModel):
     notifications_channel_id: int
diff --git a/src/services/notification.py b/src/services/notification.py
index 91931d6..895f615 100644
--- a/src/services/notification.py
+++ b/src/services/notification.py
@@ -1,8 +1,10 @@
 import logging
+from typing import Literal
 
 from httpx import AsyncClient
 
 from config import config, StreamerConfig
+from services.twitch import State
 
 
 logger = logging.getLogger(__name__)
@@ -32,11 +34,43 @@ async def notify_discord(msg: str, channel_id: str):
         )
 
 
-async def notify(msg: str, streamer_config: StreamerConfig):
+def get_role_id(streamer_config: StreamerConfig, category: str) -> int | None:
+    discord_integration = streamer_config.integrations.discord
+    if discord_integration is None:
+        return None
+
+    roles= discord_integration.roles
+    if roles is None:
+        return None
+
+    return roles.get(category)
+
+
+async def notify(notification_type: Literal["start"] | Literal["change_category"], streamer_config: StreamerConfig, current_state: State):
+    if notification_type == "start":
+        message_template = streamer_config.notifications.start_stream
+    else:
+        message_template = streamer_config.notifications.change_category
+
+    if message_template is None:
+        return
+
     integrations = streamer_config.integrations
 
     if (discord := integrations.discord) is not None:
         if discord.notifications_channel_id is not None:
+            role_id = get_role_id(streamer_config, current_state.category)
+            if role_id is not None:
+                role = f"<@&{role_id}>"
+            else:
+                role = ""
+
+            msg = message_template.format(
+                title=current_state.title,
+                category=current_state.category,
+                role=role
+            )
+
             try:
                 await notify_discord(msg, str(discord.notifications_channel_id))
             except Exception as e:
@@ -44,6 +78,12 @@ async def notify(msg: str, streamer_config: StreamerConfig):
 
     if (telegram := integrations.telegram) is not None:
         if telegram.notifications_channel_id is not None:
+            msg = message_template.format(
+                title=current_state.title,
+                category=current_state.category,
+                role=""
+            )
+
             try:
                 await notify_telegram(msg, str(telegram.notifications_channel_id))
             except Exception as e:
diff --git a/src/services/twitch.py b/src/services/twitch.py
index 6647364..4e8f385 100644
--- a/src/services/twitch.py
+++ b/src/services/twitch.py
@@ -95,12 +95,7 @@ class TwitchService:
         if streamer.notifications.start_stream is None:
             return
 
-        msg = streamer.notifications.start_stream.format(
-            title=current_state.title,
-            category=current_state.category
-        )
-
-        await notify(msg, streamer)
+        await notify("start", streamer, current_state)
 
     async def notify_change_category(self, streamer_id: int):
         current_state = self.state.get(streamer_id)
@@ -116,12 +111,7 @@ class TwitchService:
         if streamer.notifications.change_category is None:
             return
 
-        msg = streamer.notifications.change_category.format(
-            title=current_state.title,
-            category=current_state.category
-        )
-
-        await notify(msg, streamer)
+        await notify("change_category", streamer, current_state)
 
     async def get_current_stream(self, streamer_id: int, retry_count: int = 5, delay: int = 5):
         remain_retry = retry_count