bot.py

import discord
from discord.ext import commands
from discord import app_commands
from discord.ui import View, Select
from collections import defaultdict
import sqlite3
import re
import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv(".env")

# Get the bot token from the environment variable
TOKEN = os.getenv("DISCORD_TOKEN")
if not TOKEN:
    raise ValueError("No DISCORD_TOKEN found in environment variables. Please set it in the .env file.")

# Create the bot instance with the required intents
intents = discord.Intents.default()
intents.members = True  # Allow the bot to manage members
intents.messages = True  # Enable message tracking
intents.message_content = True  # Enable content tracking
bot = commands.Bot(command_prefix="!", intents=intents)

# Database setup for word counts
WORD_DB_FILE = "word_counts.db"

# Initialize the SQLite database for word counts
def init_word_db():
    with sqlite3.connect(WORD_DB_FILE) as conn:
        cursor = conn.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS global_word_counts (
                word TEXT PRIMARY KEY,
                count INTEGER
            )
        """)
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS user_word_counts (
                user_id TEXT,
                word TEXT,
                count INTEGER,
                PRIMARY KEY (user_id, word)
            )
        """)
        conn.commit()

# Update global word counts in the database
def update_global_word_counts(words):
    with sqlite3.connect(WORD_DB_FILE) as conn:
        cursor = conn.cursor()
        for word in words:
            cursor.execute("""
                INSERT INTO global_word_counts (word, count)
                VALUES (?, 1)
                ON CONFLICT(word) DO UPDATE SET count = count + 1
            """, (word,))
        conn.commit()

# Update user-specific word counts in the database
def update_user_word_counts(user_id, words):
    with sqlite3.connect(WORD_DB_FILE) as conn:
        cursor = conn.cursor()
        for word in words:
            cursor.execute("""
                INSERT INTO user_word_counts (user_id, word, count)
                VALUES (?, ?, 1)
                ON CONFLICT(user_id, word) DO UPDATE SET count = count + 1
            """, (user_id, word))
        conn.commit()

# Fetch top global words
def fetch_top_global_words(limit):
    with sqlite3.connect(WORD_DB_FILE) as conn:
        cursor = conn.cursor()
        cursor.execute("""
            SELECT word, count FROM global_word_counts
            ORDER BY count DESC
            LIMIT ?
        """, (limit,))
        results = cursor.fetchall()
    return results

# Fetch top words for a user
def fetch_top_user_words(user_id, limit):
    with sqlite3.connect(WORD_DB_FILE) as conn:
        cursor = conn.cursor()
        cursor.execute("""
            SELECT word, count FROM user_word_counts
            WHERE user_id = ?
            ORDER BY count DESC
            LIMIT ?
        """, (user_id, limit))
        results = cursor.fetchall()
    return results

# Function to clean and extract words from a message
def extract_words(message_content):
    return re.findall(r'\w+', message_content.lower())

# Initialize the SQLite database for bot data
DB_FILE = "bot_data.db"

def init_db():
    with sqlite3.connect(DB_FILE) as conn:
        cursor = conn.cursor()

        # Create user_role table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS user_role (
                discord_id TEXT PRIMARY KEY,
                role TEXT
            )
        """)

        # Create welcome_channels table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS welcome_channels (
                guild_id TEXT PRIMARY KEY,
                channel_id TEXT
            )
        """)
        conn.commit()

def save_welcome_channel(guild_id, channel_id):
    try:
        with sqlite3.connect(DB_FILE) as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO welcome_channels (guild_id, channel_id)
                VALUES (?, ?)
                ON CONFLICT(guild_id) DO UPDATE SET channel_id=excluded.channel_id
            """, (guild_id, channel_id))
            conn.commit()
    except sqlite3.Error as e:
        print(f"Database error: {e}")

def get_welcome_channel(guild_id):
    try:
        with sqlite3.connect(DB_FILE) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT channel_id FROM welcome_channels WHERE guild_id = ?", (guild_id,))
            result = cursor.fetchone()
            return int(result[0]) if result else None
    except sqlite3.Error as e:
        print(f"Database error: {e}")

def save_user_role(discord_id, role):
    try:
        with sqlite3.connect(DB_FILE) as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO user_role (discord_id, role)
                VALUES (?, ?)
                ON CONFLICT(discord_id) DO UPDATE SET role=excluded.role
            """, (discord_id, role))
            conn.commit()
    except sqlite3.Error as e:
        print(f"Database error: {e}")

def get_user_role(discord_id):
    try:
        with sqlite3.connect(DB_FILE) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT role FROM user_role WHERE discord_id = ?", (discord_id,))
            result = cursor.fetchone()
            return result[0] if result else None
    except sqlite3.Error as e:
        print(f"Database error: {e}")

# Event: Member join
@bot.event
async def on_member_join(member: discord.Member):
    channel_id = get_welcome_channel(str(member.guild.id))
    if channel_id:
        welcome_channel = member.guild.get_channel(channel_id)
        if welcome_channel:
            await welcome_channel.send(
                f"🌟 Welcome to {member.guild.name}, {member.mention}! 🌟\n"
                "We're glad to have you here! 🎉\n"
                "Enjoy your stay and explore the server!"
            )
            await member.send(f"Welcome {member.name} to {member.guild.name}! 🎉\nHave fun exploring the server!")

# Command: Set welcome channel
@bot.tree.command(name="setwelcomechannel", description="Set the channel for welcome messages.")
@app_commands.checks.has_permissions(administrator=True)
async def set_welcome_channel(interaction: discord.Interaction, channel: discord.TextChannel):
    save_welcome_channel(str(interaction.guild.id), str(channel.id))
    await interaction.response.send_message(
        f"Welcome channel has been set to {channel.mention}.", ephemeral=True
    )

# Command to show the most used words in the server
@bot.tree.command(name="topwords", description="Shows the most used words in the server")
async def topwords(interaction: discord.Interaction, limit: int = 10):
    top_words = fetch_top_global_words(limit)
    if not top_words:
        await interaction.response.send_message("No words have been counted yet.", ephemeral=True)
        return

    response = '\n'.join([f'{word}: {count}' for word, count in top_words])
    await interaction.response.send_message(f"Top {limit} words in this server:\n{response}")

# Command to show the most used words by a user
@bot.tree.command(name="topuserwords", description="Shows the most used words by a user")
async def topuserwords(interaction: discord.Interaction, user: discord.User, limit: int = 10):
    top_words = fetch_top_user_words(str(user.id), limit)
    if not top_words:
        await interaction.response.send_message(f"{user.name} has not used any words yet.", ephemeral=True)
        return

    response = '\n'.join([f'{word}: {count}' for word, count in top_words])
    await interaction.response.send_message(f"Top {limit} words used by {user.name}:\n{response}")

# Role Select class
class RoleSelect(Select):
    def __init__(self, roles):
        options = [
            discord.SelectOption(label=role.name, value=str(role.id))
            for role in roles
            if not role.is_default()  # Exclude default roles like @everyone
        ]
        super().__init__(placeholder="Choose a role", min_values=1, max_values=1, options=options)

    async def callback(self, interaction: discord.Interaction):
        role_id = int(self.values[0])
        role = interaction.guild.get_role(role_id)
        user = interaction.user

        if role:
            try:
                save_user_role(str(user.id), role.name)

                if role in user.roles:
                    await interaction.response.send_message(f"You already have the role `{role.name}`.", ephemeral=True)
                else:
                    await user.add_roles(role)
                    await interaction.response.send_message(f"Role `{role.name}` has been assigned to you.", ephemeral=True)
            except discord.Forbidden:
                await interaction.response.send_message(
                    "I don't have permission to assign roles. Please contact an admin.", ephemeral=True
                )
            except discord.HTTPException as e:
                await interaction.response.send_message(f"An error occurred: {str(e)}.", ephemeral=True)
            except Exception as e:
                print(f"Error assigning role: {e}")
                await interaction.response.send_message(
                    "An error occurred while assigning the role. Please try again later.", ephemeral=True
                )
        else:
            await interaction.response.send_message("The selected role is invalid.", ephemeral=True)

# Role Select View
class RoleSelectView(View):
    def __init__(self, roles):
        super().__init__(timeout=180)  # Set timeout to 3 minutes
        self.add_item(RoleSelect(roles))

# Command: Select role
@bot.tree.command(name="select-role", description="Select your role")
async def select_role(interaction: discord.Interaction):
    roles = interaction.guild.roles
    view = RoleSelectView(roles)
    await interaction.channel.send("Select a role from the menu below:", view=view)

# Event: on_message to update word counts
@bot.event
async def on_message(message):
    if message.author.bot:
        return

    words = extract_words(message.content)

    if words:
        update_global_word_counts(words)
        update_user_word_counts(str(message.author.id), words)

    # Ensure commands are processed after word counting
    await bot.process_commands(message)

# Initialize databases
init_word_db()
init_db()

# Event: on_ready to sync commands
@bot.event
async def on_ready():
    await bot.tree.sync()
    print(f"Bot is ready and commands are synced!")

# Run the bot
bot.run(TOKEN)