feat(roleplay_limiter): Use surge-based removal model

3 files changed, 79 insertions, 17 deletions

diff --git a/src/umabot/bot.py b/src/umabot/bot.py
index 13f3a84..6c9797c 100644
--- a/src/umabot/bot.py
+++ b/src/umabot/bot.py
@@ -75,7 +75,7 @@ class UmaBot:
         # Initialize rules
         self.rules = [
             SpamDetector(config),
-            RoleplayLimiter(config)
+            RoleplayLimiter(config, self.subreddit)
         ]
         
         # Track processed submissions to avoid processing old posts
diff --git a/src/umabot/config.py b/src/umabot/config.py
index f1cdf92..8748693 100644
--- a/src/umabot/config.py
+++ b/src/umabot/config.py
@@ -34,6 +34,11 @@ class Config:
     roleplay_limit_window_hours: int = 24  # hours
     dry_run: bool = False
     
+    # Surge-based roleplay limiting
+    roleplay_surge_threshold_1: int = 20  # First threshold for surge detection
+    roleplay_surge_threshold_2: int = 40  # Second threshold for surge detection
+    roleplay_surge_threshold_3: int = 60  # Third threshold for surge detection
+    
     @classmethod
     def from_env(cls) -> "Config":
         """Create configuration from environment variables."""
@@ -58,6 +63,9 @@ class Config:
             post_limit_window_hours=int(os.getenv("POST_LIMIT_WINDOW_HOURS", "24")),
             roleplay_limit_window_hours=int(os.getenv("ROLEPLAY_LIMIT_WINDOW_HOURS", "24")),
             dry_run=os.getenv("DRY_RUN", "false").lower() == "true",
+            roleplay_surge_threshold_1=int(os.getenv("ROLEPLAY_SURGE_THRESHOLD_1", "20")),
+            roleplay_surge_threshold_2=int(os.getenv("ROLEPLAY_SURGE_THRESHOLD_2", "40")),
+            roleplay_surge_threshold_3=int(os.getenv("ROLEPLAY_SURGE_THRESHOLD_3", "60")),
         )
     
     def validate(self) -> None:
diff --git a/src/umabot/rules/roleplay_limiter.py b/src/umabot/rules/roleplay_limiter.py
index f7724e2..60490a9 100644
--- a/src/umabot/rules/roleplay_limiter.py
+++ b/src/umabot/rules/roleplay_limiter.py
@@ -1,4 +1,4 @@
-"""Roleplay post limiter rule."""
+"""Surge-based roleplay post limiter rule."""
 
 import time
 from typing import Dict, List
@@ -7,18 +7,26 @@ from .base import Rule
 
 
 class RoleplayLimiter(Rule):
-    """Limits users to a configurable number of roleplay posts per day."""
+    """Surge-based roleplay limiter that adjusts limits based on subreddit activity."""
     
-    def __init__(self, config):
-        """Initialize the roleplay limiter."""
+    def __init__(self, config, subreddit):
+        """Initialize the surge-based roleplay limiter."""
         super().__init__(config)
+        self.subreddit = subreddit
         self.user_roleplay_posts: Dict[str, List[float]] = {}
-        self.max_roleplay_posts = config.max_roleplay_posts_per_day
-        self.time_window = config.roleplay_limit_window_hours * 60 * 60  # Convert hours to seconds
+        self.surge_window = config.roleplay_limit_window_hours * 60 * 60  # Convert hours to seconds
         self.roleplay_flair = "Roleplay"
+        
+        # Surge-based limits configuration
+        self.base_limit = 5  # Base limit when no surge (maximum posts per user)
+        self.surge_thresholds = [
+            (config.roleplay_surge_threshold_1, 3),   # First threshold = 3 posts per user
+            (config.roleplay_surge_threshold_2, 1),   # Second threshold = 1 post per user  
+            (config.roleplay_surge_threshold_3, 0),   # Third threshold = 0 posts per user (block all)
+        ]
     
     def should_remove(self, submission: praw.models.Submission) -> bool:
-        """Check if a user has posted too many roleplay posts."""
+        """Check if a user has posted too many roleplay posts based on surge."""
         if not submission.author:
             return False
         
@@ -29,38 +37,84 @@ class RoleplayLimiter(Rule):
         username = submission.author.name
         current_time = time.time()
         
+        # Get current surge level and user limit
+        surge_level, user_limit = self._get_surge_level_and_limit()
+        
         # Clean old posts from tracking
         self._clean_old_posts(username, current_time)
         
-        # Count current roleplay posts in the time window
+        # Count current roleplay posts for this user
         if username not in self.user_roleplay_posts:
             self.user_roleplay_posts[username] = []
         
-        post_count = len(self.user_roleplay_posts[username])
+        user_post_count = len(self.user_roleplay_posts[username])
         
         # Add current post to tracking
         self.user_roleplay_posts[username].append(current_time)
         
-        # Check if this post exceeds the limit
-        if post_count >= self.max_roleplay_posts:
+        # Check if this post exceeds the user's limit
+        if user_post_count >= user_limit:
             self.logger.info(
-                f"User {username} has posted {post_count + 1} roleplay posts in {self.config.roleplay_limit_window_hours} hours "
-                f"(limit: {self.max_roleplay_posts})"
+                f"User {username} has posted {user_post_count + 1} roleplay posts "
+                f"(limit: {user_limit} due to surge level: {surge_level} roleplay posts in {self.config.roleplay_limit_window_hours}h)"
             )
             return True
         
         return False
     
     def get_removal_message(self, submission: praw.models.Submission) -> str:
-        """Get the roleplay removal message."""
-        return self.config.roleplay_message
+        """Get the dynamic roleplay removal message."""
+        surge_level, user_limit = self._get_surge_level_and_limit()
+        
+        if user_limit == 0:
+            return (
+                f"Your post has been removed. Due to high roleplay activity "
+                f"({surge_level} roleplay posts in the last {self.config.roleplay_limit_window_hours} hours), "
+                f"roleplay posts are temporarily restricted in r/{self.config.subreddit_name}."
+            )
+        else:
+            return (
+                f"Your post has been removed. Due to high roleplay activity "
+                f"({surge_level} roleplay posts in the last {self.config.roleplay_limit_window_hours} hours), "
+                f"users in r/{self.config.subreddit_name} can submit {user_limit} roleplay post(s) "
+                f"within a {self.config.roleplay_limit_window_hours}-hour time window."
+            )
+    
+    def _get_surge_level_and_limit(self) -> tuple[int, int]:
+        """Get current surge level and corresponding user limit."""
+        try:
+            # Get recent roleplay posts from subreddit
+            current_time = time.time()
+            cutoff_time = current_time - self.surge_window
+            
+            # Count roleplay posts in the surge window
+            surge_count = 0
+            for submission in self.subreddit.new(limit=100):
+                if submission.created_utc < cutoff_time:
+                    break
+                if self._is_roleplay_post(submission):
+                    surge_count += 1
+            
+            # Determine limit based on surge level
+            user_limit = self.base_limit
+            for threshold, limit in self.surge_thresholds:
+                if surge_count >= threshold:
+                    user_limit = limit
+                else:
+                    break
+            
+            return surge_count, user_limit
+            
+        except Exception as e:
+            self.logger.error(f"Error calculating surge level: {e}")
+            return 0, self.base_limit
     
     def _clean_old_posts(self, username: str, current_time: float) -> None:
         """Remove roleplay posts older than the time window from tracking."""
         if username not in self.user_roleplay_posts:
             return
         
-        cutoff_time = current_time - self.time_window
+        cutoff_time = current_time - self.surge_window
         self.user_roleplay_posts[username] = [
             post_time for post_time in self.user_roleplay_posts[username]
             if post_time > cutoff_time