252 lines
9.0 KiB
Python
252 lines
9.0 KiB
Python
import discord
|
|
from discord.ext import commands
|
|
from discord import app_commands
|
|
import os
|
|
from dotenv import load_dotenv
|
|
import asyncio
|
|
import logging
|
|
import asyncpg
|
|
import redis.asyncio as aioredis
|
|
from error_handler import handle_error
|
|
import settings_manager
|
|
from global_bot_accessor import set_bot_instance
|
|
|
|
# Load environment variables from .env file
|
|
load_dotenv()
|
|
|
|
# --- Constants ---
|
|
DEFAULT_PREFIX = "!"
|
|
CORE_COGS = {"SettingsCog", "HelpCog"} # Cogs that cannot be disabled
|
|
|
|
|
|
# --- Dynamic Prefix Function ---
|
|
async def get_prefix(bot_instance, message):
|
|
"""Determines the command prefix based on guild settings or default, but disables mention as prefix."""
|
|
if not message.guild:
|
|
# Use default prefix in DMs
|
|
return DEFAULT_PREFIX
|
|
|
|
# Fetch prefix from settings manager (cache first, then DB)
|
|
prefix = await settings_manager.get_guild_prefix(message.guild.id, DEFAULT_PREFIX)
|
|
return prefix
|
|
|
|
|
|
# --- Bot Setup ---
|
|
# Set up intents (permissions)
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
intents.members = True
|
|
|
|
|
|
# --- Custom Bot Class with setup_hook for async initialization ---
|
|
class NeruBot(commands.Bot):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
owner_id = os.getenv("OWNER_USER_ID")
|
|
if owner_id:
|
|
self.owner_id = int(owner_id)
|
|
self.core_cogs = CORE_COGS # Attach core cogs list to bot instance
|
|
self.settings_manager = settings_manager # Attach settings manager instance
|
|
self.pg_pool = None # Will be initialized in setup_hook
|
|
self.redis = None # Will be initialized in setup_hook
|
|
self.ai_cogs_to_skip = [] # For --disable-ai flag
|
|
|
|
async def setup_hook(self):
|
|
"""Async initialization that runs after login but before on_ready."""
|
|
log.info("NeruBot setup_hook called")
|
|
|
|
# Initialize database connections
|
|
try:
|
|
# PostgreSQL connection pool
|
|
self.pg_pool = await asyncpg.create_pool(
|
|
user=os.getenv("POSTGRES_USER"),
|
|
password=os.getenv("POSTGRES_PASSWORD"),
|
|
host=os.getenv("POSTGRES_HOST"),
|
|
database=os.getenv("POSTGRES_SETTINGS_DB"),
|
|
)
|
|
log.info("PostgreSQL connection pool initialized")
|
|
|
|
# Redis connection
|
|
self.redis = await aioredis.from_url(
|
|
f"redis://{os.getenv('REDIS_HOST')}:{os.getenv('REDIS_PORT', '6379')}",
|
|
password=os.getenv("REDIS_PASSWORD"),
|
|
decode_responses=True,
|
|
)
|
|
log.info("Redis connection initialized")
|
|
|
|
# Initialize database schema and run migrations using settings_manager
|
|
if self.pg_pool and self.redis:
|
|
try:
|
|
await settings_manager.initialize_database()
|
|
log.info(
|
|
"Database schema initialization called via settings_manager."
|
|
)
|
|
await settings_manager.run_migrations()
|
|
log.info("Database migrations called via settings_manager.")
|
|
except Exception as e:
|
|
log.exception(
|
|
"CRITICAL: Failed during settings_manager database setup (init/migrations)."
|
|
)
|
|
else:
|
|
log.error(
|
|
"CRITICAL: pg_pool or redis_client not initialized in setup_hook. Cannot proceed with settings_manager setup."
|
|
)
|
|
|
|
# Load only specific cogs
|
|
try:
|
|
# Define the hardcoded list of cogs to load
|
|
hardcoded_cogs = [
|
|
"cogs.settings_cog",
|
|
"cogs.help_cog",
|
|
"cogs.neru_message_cog",
|
|
"cogs.owoify_cog",
|
|
"cogs.caption_cog",
|
|
"cogs.games_cog",
|
|
"cogs.ping_cog",
|
|
"cogs.neru_teto_cog",
|
|
"cogs.random_strings_cog",
|
|
"cogs.neru_roleplay_cog",
|
|
"cogs.system_check_cog",
|
|
"cogs.teto_image_cog",
|
|
"cogs.webdrivertorso_cog",
|
|
"cogs.ban_system_cog",
|
|
"cogs.terminal_cog",
|
|
"cogs.shell_command_cog",
|
|
"cogs.marriage_cog",
|
|
"cogs.upload_cog",
|
|
"cogs.dictionary_cog",
|
|
"cogs.eval_cog",
|
|
"cogs.fetch_user_cog",
|
|
]
|
|
|
|
# Load each cog individually
|
|
for cog in hardcoded_cogs:
|
|
try:
|
|
await self.load_extension(cog)
|
|
log.info(f"Successfully loaded {cog}")
|
|
except Exception as e:
|
|
log.exception(f"Error loading cog {cog}: {e}")
|
|
|
|
log.info(f"Loaded {len(hardcoded_cogs)} hardcoded cogs")
|
|
except Exception as e:
|
|
log.exception(f"Error during cog loading process: {e}")
|
|
|
|
# Apply global allowed_installs and allowed_contexts to all commands
|
|
try:
|
|
log.info(
|
|
"Applying global allowed_installs and allowed_contexts to all commands..."
|
|
)
|
|
for command in self.tree.get_commands():
|
|
# Apply decorators to each command
|
|
app_commands.allowed_installs(guilds=True, users=True)(command)
|
|
app_commands.allowed_contexts(
|
|
guilds=True, dms=True, private_channels=True
|
|
)(command)
|
|
|
|
# Sync commands globally
|
|
log.info("Starting global command sync process...")
|
|
synced = await self.tree.sync()
|
|
log.info(f"Synced {len(synced)} commands globally")
|
|
|
|
# List commands after sync
|
|
commands_after = [cmd.name for cmd in self.tree.get_commands()]
|
|
log.info(f"Commands registered in command tree: {commands_after}")
|
|
except Exception as e:
|
|
log.exception(f"Failed to sync commands: {e}")
|
|
|
|
except Exception as e:
|
|
log.exception(f"Error in setup_hook: {e}")
|
|
|
|
|
|
# Create bot instance using the custom class
|
|
bot = NeruBot(command_prefix=get_prefix, intents=intents)
|
|
|
|
# --- Logging Setup ---
|
|
# Configure logging (adjust level and format as needed)
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s:%(levelname)s:%(name)s: %(message)s"
|
|
)
|
|
log = logging.getLogger(__name__) # Logger for neru_bot.py
|
|
|
|
|
|
# --- Events ---
|
|
@bot.event
|
|
async def on_ready():
|
|
if bot.user:
|
|
log.info(f"{bot.user.name} has connected to Discord!")
|
|
log.info(f"Bot ID: {bot.user.id}")
|
|
# Set the bot's status
|
|
await bot.change_presence(
|
|
activity=discord.Activity(type=discord.ActivityType.listening, name="!help")
|
|
)
|
|
log.info("Bot status set to 'Listening to !help'")
|
|
|
|
|
|
# Error handling
|
|
@bot.event
|
|
async def on_command_error(ctx, error):
|
|
await handle_error(ctx, error)
|
|
|
|
|
|
@bot.tree.error
|
|
async def on_app_command_error(interaction, error):
|
|
await handle_error(interaction, error)
|
|
|
|
|
|
async def main(args):
|
|
"""Main async function to load cogs and start the bot."""
|
|
TOKEN = os.getenv("NERU_BOT_TOKEN")
|
|
if not TOKEN:
|
|
raise ValueError(
|
|
"No token found. Make sure to set NERU_BOT_TOKEN in your .env file."
|
|
)
|
|
|
|
# Set the global bot instance for other modules to access
|
|
set_bot_instance(bot)
|
|
log.info(f"Global bot instance set in global_bot_accessor. Bot ID: {id(bot)}")
|
|
|
|
# Configure AI cogs to skip if needed
|
|
if args.disable_ai:
|
|
log.info("AI functionality disabled via command line flag.")
|
|
ai_cogs_to_skip = [
|
|
"cogs.ai_cog",
|
|
"cogs.multi_conversation_ai_cog",
|
|
# Add any other AI-related cogs from the 'cogs' folder here
|
|
]
|
|
# Store the skip list on the bot object for reload commands
|
|
bot.ai_cogs_to_skip = ai_cogs_to_skip
|
|
else:
|
|
bot.ai_cogs_to_skip = [] # Ensure it exists even if empty
|
|
|
|
try:
|
|
# The bot will call setup_hook internally after login but before on_ready.
|
|
await bot.start(TOKEN)
|
|
except Exception as e:
|
|
log.exception(f"An error occurred during bot.start(): {e}")
|
|
finally:
|
|
# Close database connections
|
|
if bot.pg_pool:
|
|
await bot.pg_pool.close()
|
|
log.info("PostgreSQL connection pool closed.")
|
|
if bot.redis:
|
|
await bot.redis.close()
|
|
log.info("Redis connection closed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
# Parse command line arguments
|
|
parser = argparse.ArgumentParser(description="Run the Neru Discord Bot.")
|
|
parser.add_argument(
|
|
"--disable-ai", action="store_true", help="Disable AI functionality"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
asyncio.run(main(args))
|
|
except KeyboardInterrupt:
|
|
log.info("Bot stopped by user.")
|
|
except Exception as e:
|
|
log.exception(f"An error occurred running the bot: {e}")
|