titlebot-ng-ng/get_user_list.py

"""
Python script that reads configs/config.yaml and scrapes a list of active users from the Matrix room.
Stores the user list in a SQLite database named after the room.
"""

import yaml
import re
import asyncio
import sqlite3
import time
import sys
from nio import AsyncClient, LoginResponse, RoomMessageText, RoomMessagesResponse, SyncResponse
from room import Room
import dbutil


class UserListManager:
    """Manages collection and storage of active users in a Matrix room."""
    
    def __init__(self, config_path="configs/config.yaml"):
        self.room = Room(config_path)
        self.users = []
        self.db_path = self.get_user_db_path()
        self.init_user_db()
    
    def get_user_db_path(self):
        """Generate database path for user list, similar to suggestions database."""
        safe = re.sub(r'[^a-zA-Z0-9_-]', '_', self.room.config.room_id)
        return f"users_{safe}.db"
    
    def init_user_db(self):
        """Initialize the user list database."""
        conn = dbutil.connect(self.db_path)
        c = conn.cursor()
        c.execute('''
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id TEXT UNIQUE,
                display_name TEXT,
                last_seen INTEGER,
                last_updated INTEGER,
                num_posts INTEGER DEFAULT 0,
                last_active INTEGER DEFAULT 0,
                UNIQUE(user_id)
            )
        ''')
        # Add num_posts column if it doesn't exist (for existing databases)
        c.execute("PRAGMA table_info(users)")
        cols = [row[1] for row in c.fetchall()]
        if 'num_posts' not in cols:
            c.execute('ALTER TABLE users ADD COLUMN num_posts INTEGER DEFAULT 0')
        if 'last_active' not in cols:
            c.execute('ALTER TABLE users ADD COLUMN last_active INTEGER DEFAULT 0')
        
        # Table to track when user list was last fetched
        c.execute('''
            CREATE TABLE IF NOT EXISTS fetch_log (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp INTEGER,
                user_count INTEGER,
                datetime TEXT,
                most_recent_post INTEGER DEFAULT 0
            )
        ''')
        # Add most_recent_post column if it doesn't exist
        c.execute("PRAGMA table_info(fetch_log)")
        cols = [row[1] for row in c.fetchall()]
        if 'most_recent_post' not in cols:
            c.execute('ALTER TABLE fetch_log ADD COLUMN most_recent_post INTEGER DEFAULT 0')
        
        dbutil.commit(conn)
        conn.close()
    
    async def get_room_members(self):
        """Fetch all members from the Matrix room and count their posts."""
        # Ensure we're logged in
        if not await self.room.login_to_matrix():
            print("Failed to login to Matrix")
            return []
        
        # Sync to get room state
        assert self.room.client is not None
        sync_resp = await self.room.client.sync(timeout=3000)
        if not isinstance(sync_resp, SyncResponse):
            print(f"Sync failed: {sync_resp}")
        
        # Get the room object
        room_obj = self.room.client.rooms.get(self.room.config.room_id)
        if not room_obj:
            print(f"Could not find room {self.room.config.room_id}")
            return []
        
        # Count posts by fetching message history
        print("Counting posts per user (this may take a moment)...")
        post_counts, last_activity, most_recent_post = await self.count_user_posts()
        
        users = []
        for user_id in room_obj.users:
            display_name = await self.room.get_display_name(user_id)
            users.append({
                'user_id': user_id,
                'display_name': display_name,
                'last_seen': int(time.time()),
                'num_posts': post_counts.get(user_id, 0),
                'last_active': last_activity.get(user_id, 0)
            })
        
        self.users = users
        self.most_recent_post = most_recent_post
        return users
    
    async def count_user_posts(self):
        """Count total number of posts by each user in the room and track last activity."""
        post_counts = {}
        last_activity = {}
        most_recent_post = 0
        batch_size = self.room.config.batch_size
        lookback_limit = self.room.config.user_lookback_limit
        next_batch = None
        fetched = 0
        
        # Fetch messages in batches to count posts per user
        assert self.room.client is not None
        while fetched < lookback_limit:
            try:
                response = await self.room.client.room_messages(
                    room_id=self.room.config.room_id,
                    start=next_batch,
                    limit=batch_size
                )
                
                if not isinstance(response, RoomMessagesResponse) or not response.chunk:
                    break
                
                fetched += len(response.chunk)
                
                for event in response.chunk:
                    if isinstance(event, RoomMessageText):
                        sender = event.sender
                        post_counts[sender] = post_counts.get(sender, 0) + 1
                        
                        # Track the most recent message timestamp for each user
                        event_time = int(event.server_timestamp // 1000)
                        if sender not in last_activity or event_time > last_activity[sender]:
                            last_activity[sender] = event_time
                        # Track the absolute most recent post in the room
                        if event_time > most_recent_post:
                            most_recent_post = event_time
                
                next_batch = response.end
                
                # Stop if we've reached the end or the lookback limit
                if not next_batch or fetched >= lookback_limit:
                    break
                    
            except Exception as e:
                print(f"Error fetching messages: {e}")
                break
        
        return post_counts, last_activity, most_recent_post
    
    def write_users_to_db(self):
        """Write collected users to the database."""
        if not self.users:
            print("No users to write to database")
            return
        
        conn = dbutil.connect(self.db_path)
        c = conn.cursor()
        
        current_time = int(time.time())
        
        for user in self.users:
            try:
                c.execute('''
                    INSERT INTO users (user_id, display_name, last_seen, last_updated, num_posts, last_active)
                    VALUES (?, ?, ?, ?, ?, ?)
                    ON CONFLICT(user_id) DO UPDATE SET
                        display_name = excluded.display_name,
                        last_seen = excluded.last_seen,
                        last_updated = excluded.last_updated,
                        num_posts = excluded.num_posts,
                        last_active = excluded.last_active
                ''', (
                    user['user_id'],
                    user['display_name'],
                    user['last_seen'],
                    current_time,
                    user['num_posts'],
                    user.get('last_active', 0)
                ))
            except sqlite3.IntegrityError:
                # Update existing record
                c.execute('''
                    UPDATE users 
                    SET display_name = ?, last_seen = ?, last_updated = ?, num_posts = ?, last_active = ?
                    WHERE user_id = ?
                ''', (
                    user['display_name'],
                    user['last_seen'],
                    current_time,
                    user['num_posts'],
                    user.get('last_active', 0),
                    user['user_id']
                ))
        
        # Log this fetch
        from datetime import datetime
        c.execute('''
            INSERT INTO fetch_log (timestamp, user_count, datetime, most_recent_post)
            VALUES (?, ?, ?, ?)
        ''', (current_time, len(self.users), datetime.now().isoformat(), getattr(self, 'most_recent_post', 0)))
        
        dbutil.commit(conn)
        conn.close()
        print(f"Wrote {len(self.users)} users to {self.db_path}")
    
    def read_users_from_db(self):
        """Read all users from the database."""
        conn = dbutil.connect(self.db_path)
        c = conn.cursor()
        c.execute('SELECT user_id, display_name, last_seen, last_updated, num_posts, last_active FROM users ORDER BY num_posts DESC, display_name')
        users = []
        for row in c.fetchall():
            users.append({
                'user_id': row[0],
                'display_name': row[1],
                'last_seen': row[2],
                'last_updated': row[3],
                'num_posts': row[4],
                'last_active': row[5] if len(row) > 5 else 0
            })
        conn.close()
        return users
    
    async def close(self):
        """Close Matrix client connection."""
        await self.room.close()


async def main():
    """Main function to fetch and store user list continuously."""
    config_path = sys.argv[1] if len(sys.argv) > 1 else "configs/config.yaml"
    manager = UserListManager(config_path)
    
    try:
        print("Starting user list collection service...")
        print("This will run continuously, updating the user list every 5 minutes.")
        
        while True:
            try:
                print("\nFetching user list from Matrix room...")
                users = await manager.get_room_members()
                
                if users:
                    print(f"Found {len(users)} users in room")
                    manager.write_users_to_db()
                    
                    # Display top users by post count
                    print("\nTop users by post count:")
                    sorted_users = sorted(users, key=lambda u: u['num_posts'], reverse=True)
                    for user in sorted_users[:10]:  # Show top 10
                        print(f"  - {user['display_name']}: {user['num_posts']} posts")
                    
                    if len(sorted_users) > 10:
                        print(f"  ... and {len(sorted_users) - 10} more users")
                else:
                    print("No users found or could not connect to room")
                
                # Wait 5 minutes before next update
                print("\nWaiting 5 minutes before next update...")
                await asyncio.sleep(300)  # 5 minutes
                
            except Exception as e:
                print(f"Error during user list update: {e}")
                print("Retrying in 1 minute...")
                await asyncio.sleep(60)  # Wait 1 minute on error before retry
    
    except KeyboardInterrupt:
        print("\nShutting down user list service...")
    finally:
        await manager.close()


if __name__ == "__main__":
    asyncio.run(main())