Add voting system for !traps command (3+ votes, 15 min timeout)

Co-authored-by: kuhyx <147418882+kuhyx@users.noreply.github.com>

main.py

@@ -6,6 +6,7 @@ import base64
 import json
 from datetime import datetime, time, timedelta
 from fastapi import FastAPI
+from rule34Py import rule34Py
 
 # Create FastAPI app
 app = FastAPI()
@@ -20,6 +21,9 @@ CAT_API = os.getenv('CAT_API', '')
 last_command_time = None
 warning_sent = False
 
+# Initialize rule34Py for trap images
+r34_client = rule34Py()
+
 class StringCounter:
     def __init__(self):
         self.string_map = {}
@@ -36,6 +40,60 @@ class StringCounter:
             return self.string_map[key]['common_name']
         return None
 
+
+class VotingSystem:
+    """Tracks votes for commands that require multiple users to agree."""
+    REQUIRED_VOTES = 3
+    VOTE_TIMEOUT_MINUTES = 15
+
+    def __init__(self):
+        self.votes = {}  # command -> {user_uuid: timestamp}
+
+    def add_vote(self, command, user_uuid):
+        """Add a vote from a user. Returns (vote_count, newly_passed)."""
+        current_time = datetime.now()
+        if command not in self.votes:
+            self.votes[command] = {}
+
+        # Remove expired votes
+        self._cleanup_expired_votes(command, current_time)
+
+        # Add or update the user's vote
+        already_voted = user_uuid in self.votes[command]
+        self.votes[command][user_uuid] = current_time
+
+        vote_count = len(self.votes[command])
+        # Check if we just reached the threshold
+        newly_passed = vote_count >= self.REQUIRED_VOTES and not already_voted
+
+        return vote_count, newly_passed
+
+    def _cleanup_expired_votes(self, command, current_time):
+        """Remove votes older than VOTE_TIMEOUT_MINUTES."""
+        if command not in self.votes:
+            return
+        timeout = timedelta(minutes=self.VOTE_TIMEOUT_MINUTES)
+        self.votes[command] = {
+            user_uuid: timestamp
+            for user_uuid, timestamp in self.votes[command].items()
+            if current_time - timestamp < timeout
+        }
+
+    def reset_votes(self, command):
+        """Reset votes for a command after it has been triggered."""
+        if command in self.votes:
+            self.votes[command] = {}
+
+    def get_vote_count(self, command):
+        """Get current vote count for a command."""
+        current_time = datetime.now()
+        self._cleanup_expired_votes(command, current_time)
+        return len(self.votes.get(command, {}))
+
+
+# Global voting system instance
+voting_system = VotingSystem()
+
 def download_image(image_url):
     # Download the image
     image_response = requests.get(image_url)
@@ -121,10 +179,25 @@ def extract_source_uuid(message):
 
 command_map = {
     ("!kot", "!koty", "!kots", "!cat", "!cats", "!meow", "!miau", "!ᴋᴏᴛ", "!𝓴𝓸𝓽", "!𝗸𝗼𝘁"): lambda recipient: send_image(fetch_and_download_image("https://api.thecatapi.com/v1/images/search", [0, 'url']), recipient),
-    ("!pies", "!psy", "!dog", "!dogs", "!woof", "!szczek", "!𝗽𝗶𝗲𝘀", "!͓̽p͓̽i͓̽e͓̽s͓̽"): lambda recipient: send_image(fetch_and_download_image("https://dog.ceo/api/breeds/image/random", 'message'), recipient),
-    # ("!traps"): lambda recipient: send_image(download_image(((r34Py.random_post(["trap"])).sample)), recipient)
+    ("!pies", "!psy", "!dog", "!dogs", "!woof", "!szczek", "!𝗽𝗶𝗲𝘀", "!͓̽p͓̽i͓̽e͓̽s͓̽"): lambda recipient: send_image(fetch_and_download_image("https://dog.ceo/api/breeds/image/random", 'message'), recipient),
 }
 
+# Commands that require voting (3+ votes within 15 minutes)
+VOTING_COMMANDS = ("!traps", "!trap")
+
+
+async def send_trap_image(recipient):
+    """Fetch and send a trap image from rule34."""
+    try:
+        post = r34_client.random_post(["trap"])
+        if post and post.sample:
+            base64_data = download_image(post.sample)
+            await send_image(base64_data, recipient)
+        else:
+            send_message("Nie znaleziono obrazka.", recipient)
+    except Exception as e:
+        send_message(f"Błąd podczas pobierania obrazka: {str(e)}", recipient)
+
 def extract_source_name(message):
     message_json = message
     inside_message = message_json.get('sourceName', {})
@@ -154,13 +227,20 @@ async def scheduled_task(counter):
         send_message(counter.string_map, GROUP_ID_SEND)
         counter.string_map = {}
 
-async def trigger_command(message_content, recipient):
+async def trigger_command(message_content, recipient, user_uuid=None):
     global last_command_time, warning_sent
     message_value = message_message(message_content)
     
     try:
         if message_value is not None and message_value[0] == "!":
             current_time = datetime.now()
+
+            # Handle voting commands separately (no cooldown for voting)
+            if message_value in VOTING_COMMANDS:
+                if user_uuid:
+                    await handle_voting_command(message_value, user_uuid, recipient)
+                return
+
             if last_command_time and current_time - last_command_time < timedelta(seconds=10):
                 if not warning_sent:
                     send_message("BEEP BOOP POCZEKAJ 10 SEKUND.", recipient)
@@ -178,10 +258,46 @@ async def trigger_command(message_content, recipient):
     except Exception as e:
         send_message(f"trigger_command, unknown error {message_content}: {str(e)}", recipient)
 
+
+async def handle_voting_command(command, user_uuid, recipient):
+    """Handle commands that require voting."""
+    global last_command_time, warning_sent
+    vote_count, newly_passed = voting_system.add_vote("traps", user_uuid)
+    required = VotingSystem.REQUIRED_VOTES
+    timeout = VotingSystem.VOTE_TIMEOUT_MINUTES
+
+    if newly_passed:
+        # Check cooldown only when threshold is reached
+        current_time = datetime.now()
+        if last_command_time and current_time - last_command_time < timedelta(seconds=10):
+            if not warning_sent:
+                send_message("BEEP BOOP POCZEKAJ 10 SEKUND.", recipient)
+                warning_sent = True
+            return
+
+        send_message(f"Głosowanie zakończone! ({vote_count}/{required}) Wysyłam obrazek...", recipient)
+        await send_trap_image(recipient)
+        voting_system.reset_votes("traps")
+        last_command_time = current_time
+        warning_sent = False
+    elif vote_count >= required:
+        # Already passed in a previous vote, just inform
+        send_message("Już wysłano obrazek. Głosowanie zresetowane.", recipient)
+        voting_system.reset_votes("traps")
+    else:
+        remaining_votes = required - vote_count
+        send_message(
+            f"Głos zapisany! ({vote_count}/{required}) "
+            f"Potrzeba jeszcze {remaining_votes} głos(ów) w ciągu {timeout} minut.",
+            recipient
+        )
+
 async def send_to_group(message_content, counter, message):
     if message_group_id(message_content) == GROUP_ID:
-        await count_messages(json.loads(message).get('envelope', {}), counter)
-        await trigger_command(message_content, GROUP_ID_SEND)
+        envelope = json.loads(message).get('envelope', {})
+        await count_messages(envelope, counter)
+        user_uuid = extract_source_uuid(envelope)
+        await trigger_command(message_content, GROUP_ID_SEND, user_uuid)
 
 async def remove_attachment(attachment_id):
     response = requests.delete(REMOVE_ATTACHMENT_URL + attachment_id)