mrrrme/database/db_tool.py

"""Database Cleanup and Verification Tool for MrrrMe"""
import sqlite3
import os

DB_PATH = "/data/mrrrme_users.db" if os.path.exists("/data") else "/tmp/mrrrme_users.db"

def view_all_users():
    """Show all users and their summaries"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("""
        SELECT u.user_id, u.username, u.created_at, s.summary_text, s.updated_at
        FROM users u
        LEFT JOIN user_summaries s ON u.user_id = s.user_id
        ORDER BY u.created_at DESC
    """)
    
    print("\n" + "="*80)
    print("👥 ALL USERS")
    print("="*80)
    
    for user_id, username, created, summary, updated in cursor.fetchall():
        print(f"\n📌 {username} (ID: {user_id})")
        print(f"   Created: {created}")
        if summary:
            print(f"   Summary: {summary}")
            print(f"   Updated: {updated}")
        else:
            print(f"   Summary: (none yet)")
    
    conn.close()

def view_user_messages(username):
    """Show all messages for a specific user"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("""
        SELECT m.role, m.content, m.emotion, m.timestamp
        FROM messages m
        JOIN sessions s ON m.session_id = s.session_id
        JOIN users u ON s.user_id = u.user_id
        WHERE u.username = ?
        ORDER BY m.timestamp ASC
    """, (username,))
    
    print(f"\n" + "="*80)
    print(f"💬 MESSAGES FOR {username}")
    print("="*80)
    
    for role, content, emotion, timestamp in cursor.fetchall():
        emo_tag = f" [{emotion}]" if emotion else ""
        speaker = "User" if role == "user" else "AI"
        print(f"[{timestamp}] {speaker}{emo_tag}: {content}")
    
    conn.close()

def clear_user_summary(username):
    """Clear summary for a user (for testing)"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("SELECT user_id FROM users WHERE username = ?", (username,))
    result = cursor.fetchone()
    
    if not result:
        print(f"❌ User {username} not found")
        conn.close()
        return
    
    user_id = result[0]
    
    cursor.execute("DELETE FROM user_summaries WHERE user_id = ?", (user_id,))
    conn.commit()
    conn.close()
    
    print(f"✅ Cleared summary for {username}")

def clear_all_summaries():
    """Clear all summaries (for fresh start)"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("DELETE FROM user_summaries")
    deleted = cursor.rowcount
    
    conn.commit()
    conn.close()
    
    print(f"✅ Cleared {deleted} summaries")

def delete_user(username):
    """Delete a user and all their data"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("SELECT user_id FROM users WHERE username = ?", (username,))
    result = cursor.fetchone()
    
    if not result:
        print(f"❌ User {username} not found")
        conn.close()
        return
    
    user_id = result[0]
    
    # Delete in correct order (foreign keys)
    cursor.execute("DELETE FROM messages WHERE session_id IN (SELECT session_id FROM sessions WHERE user_id = ?)", (user_id,))
    cursor.execute("DELETE FROM user_summaries WHERE user_id = ?", (user_id,))
    cursor.execute("DELETE FROM sessions WHERE user_id = ?", (user_id,))
    cursor.execute("DELETE FROM users WHERE user_id = ?", (user_id,))
    
    conn.commit()
    conn.close()
    
    print(f"✅ Deleted user {username} and all their data")

def reset_database():
    """DANGER: Delete everything and start fresh"""
    confirm = input("⚠️  This will delete ALL users and data. Type 'YES' to confirm: ")
    
    if confirm != "YES":
        print("❌ Cancelled")
        return
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("DROP TABLE IF EXISTS messages")
    cursor.execute("DROP TABLE IF EXISTS user_summaries")
    cursor.execute("DROP TABLE IF EXISTS sessions")
    cursor.execute("DROP TABLE IF EXISTS users")
    
    conn.commit()
    conn.close()
    
    # Recreate tables
    from mrrrme.backend_server_old import init_db
    init_db()
    
    print("✅ Database reset complete")

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("\n🔧 MrrrMe Database Tool")
        print("="*80)
        print("\nUsage:")
        print("  python db_tool.py list                  # Show all users")
        print("  python db_tool.py messages <username>    # Show user's messages")
        print("  python db_tool.py clear-summary <username>  # Clear user's summary")
        print("  python db_tool.py clear-all-summaries    # Clear all summaries")
        print("  python db_tool.py delete <username>      # Delete user")
        print("  python db_tool.py reset                  # Reset entire database")
        print("\nExamples:")
        print("  python db_tool.py list")
        print("  python db_tool.py messages alice")
        print("  python db_tool.py clear-summary testuser")
        print("  python db_tool.py clear-all-summaries")
        print("\n")
        sys.exit(1)
    
    command = sys.argv[1]
    
    if command == "list":
        view_all_users()
    
    elif command == "messages" and len(sys.argv) > 2:
        view_user_messages(sys.argv[2])
    
    elif command == "clear-summary" and len(sys.argv) > 2:
        clear_user_summary(sys.argv[2])
    
    elif command == "clear-all-summaries":
        clear_all_summaries()
    
    elif command == "delete" and len(sys.argv) > 2:
        delete_user(sys.argv[2])
    
    elif command == "reset":
        reset_database()
    
    else:
        print(f"❌ Unknown command: {command}")
        print("Run without arguments to see usage")