bot.py

import asyncio
import logging
import os
import io
import time
from collections import defaultdict, deque
from typing import Optional

import discord
from discord.ext import commands
from dotenv import load_dotenv

from gemini_client import GeminiImageGenerator

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Bot configuration
DISCORD_TOKEN = os.getenv('DISCORD_TOKEN')
GEMINI_API_KEY = os.getenv('GEMINI_API_KEY')

if not DISCORD_TOKEN:
    raise ValueError("DISCORD_TOKEN environment variable is required")
if not GEMINI_API_KEY:
    raise ValueError("GEMINI_API_KEY environment variable is required")

# Rate limiting configuration
RATE_LIMIT_REQUESTS = 5  # Max requests per user
RATE_LIMIT_WINDOW = 300  # Time window in seconds (5 minutes)

class RateLimiter:
    def __init__(self):
        self.user_requests = defaultdict(deque)
    
    def is_rate_limited(self, user_id: int) -> bool:
        now = time.time()
        user_queue = self.user_requests[user_id]
        
        # Remove old requests outside the time window
        while user_queue and user_queue[0] < now - RATE_LIMIT_WINDOW:
            user_queue.popleft()
        
        # Check if user has exceeded rate limit
        if len(user_queue) >= RATE_LIMIT_REQUESTS:
            return True
        
        # Add current request timestamp
        user_queue.append(now)
        return False
    
    def get_reset_time(self, user_id: int) -> Optional[int]:
        user_queue = self.user_requests[user_id]
        if user_queue:
            return int(RATE_LIMIT_WINDOW - (time.time() - user_queue[0]))
        return None

# Initialize bot with intents
intents = discord.Intents.default()
# Note: message_content is a privileged intent that needs to be enabled in Discord Developer Portal
# For now, we'll disable it and use mentions instead of prefix commands
bot = commands.Bot(command_prefix='!', intents=intents)

# Initialize rate limiter and image generator
rate_limiter = RateLimiter()
image_generator = GeminiImageGenerator(GEMINI_API_KEY)

# Helper functions for mention-based commands
async def generate_image_from_mention(message, prompt: str):
    """Generate an image from a mention-based command."""
    
    # Check rate limiting
    if rate_limiter.is_rate_limited(message.author.id):
        reset_time = rate_limiter.get_reset_time(message.author.id)
        await message.reply(f"Rate limit exceeded. Please try again in {reset_time} seconds.")
        return
    
    # Validate prompt length
    if len(prompt) > 1000:
        await message.reply("Prompt is too long. Please keep it under 1000 characters.")
        return
    
    if len(prompt.strip()) < 3:
        await message.reply("Prompt is too short. Please provide a more detailed description.")
        return
    
    # Send initial response
    async with message.channel.typing():
        try:
            # Generate image
            logger.info(f"Generating image for user {message.author.id} with prompt: {prompt[:100]}...")
            image_data = await image_generator.generate_image_async(prompt)
            
            if not image_data:
                await message.reply("Failed to generate image. The model might be experiencing issues. Please try again later.")
                return
            
            # Create Discord file object
            image_file = discord.File(io.BytesIO(image_data), filename="generated_image.png")
            
            # Create embed for better presentation
            embed = discord.Embed(
                title="Generated Image",
                description=f"**Prompt:** {prompt}",
                color=0x00ff00
            )
            embed.set_image(url="attachment://generated_image.png")
            embed.set_footer(text=f"Generated by {message.author.display_name}")
            
            # Send the image
            await message.reply(embed=embed, file=image_file)
            logger.info(f"Successfully generated and sent image for user {message.author.id}")
            
        except Exception as e:
            logger.error(f"Error generating image: {e}")
            await message.reply("Failed to generate image. Please check your prompt and try again later.")

async def help_generate_mention(message):
    """Show help information for mention-based commands."""
    embed = discord.Embed(
        title="Image Generation Help",
        description="Generate images using Google's Gemini 2.0 Flash Preview model",
        color=0x0099ff
    )
    
    embed.add_field(
        name="Commands",
        value="@bot generate <prompt> - Generate an image\n@bot gen <prompt> - Short alias\n@bot img <prompt> - Another alias",
        inline=False
    )
    
    embed.add_field(
        name="Usage Tips",
        value="• Be descriptive in your prompts\n• Include style, mood, and details\n• Prompts should be 3-1000 characters\n• Use English for best results",
        inline=False
    )
    
    embed.add_field(
        name="Examples",
        value="@bot generate a majestic dragon flying over a fantasy castle at sunset\n@bot gen cyberpunk cityscape with neon lights and rain\n@bot img cute cat wearing a wizard hat, digital art style",
        inline=False
    )
    
    embed.add_field(
        name="Rate Limits",
        value=f"• {RATE_LIMIT_REQUESTS} requests per {RATE_LIMIT_WINDOW//60} minutes per user\n• Please be patient, image generation takes time",
        inline=False
    )
    
    await message.reply(embed=embed)

async def ping_mention(message):
    """Check bot latency from mention."""
    latency = round(bot.latency * 1000)
    await message.reply(f"Pong! Latency: {latency}ms")

async def bot_info_mention(message):
    """Show bot information from mention."""
    embed = discord.Embed(
        title="Bot Information",
        color=0x0099ff
    )
    
    embed.add_field(name="Servers", value=len(bot.guilds), inline=True)
    embed.add_field(name="Users", value=len(set(bot.get_all_members())), inline=True)
    embed.add_field(name="Latency", value=f"{round(bot.latency * 1000)}ms", inline=True)
    
    embed.add_field(
        name="Features",
        value="• AI Image Generation\n• Rate Limiting\n• Error Handling\n• Mention-based Commands",
        inline=False
    )
    
    embed.set_footer(text="Powered by Google Gemini 2.0 Flash Preview")
    
    await message.reply(embed=embed)

@bot.event
async def on_ready():
    logger.info(f'{bot.user} has connected to Discord!')
    logger.info(f'Bot is in {len(bot.guilds)} guilds')

@bot.event
async def on_message(message):
    # Ignore messages from the bot itself
    if message.author == bot.user:
        return
    
    # Check if bot is mentioned
    if bot.user and bot.user.mentioned_in(message):
        content = message.content.replace(f'<@{bot.user.id}>', '').replace(f'<@!{bot.user.id}>', '').strip()
        
        if content.lower().startswith(('generate', 'gen', 'img')):
            # Extract prompt after the command
            parts = content.split(' ', 1)
            if len(parts) > 1:
                prompt = parts[1].strip()
                if prompt:
                    await generate_image_from_mention(message, prompt)
                else:
                    await message.reply("Please provide a prompt for image generation. Example: `@bot generate a beautiful sunset over mountains`")
            else:
                await message.reply("Please provide a prompt for image generation. Example: `@bot generate a beautiful sunset over mountains`")
        
        elif content.lower() in ['help', 'help generate', 'genhelp']:
            await help_generate_mention(message)
        
        elif content.lower() == 'ping':
            await ping_mention(message)
        
        elif content.lower() == 'info':
            await bot_info_mention(message)
        
        else:
            await message.reply("Hi! Mention me with `generate <prompt>` to create images, or `help` for more info.")
    
    # Process commands normally (for servers that have message content intent enabled)
    await bot.process_commands(message)

@bot.event
async def on_command_error(ctx, error):
    if isinstance(error, commands.CommandNotFound):
        return
    elif isinstance(error, commands.MissingRequiredArgument):
        await ctx.send("Please provide a prompt for image generation. Example: `!generate a beautiful sunset over mountains`")
    else:
        logger.error(f"Command error: {error}")
        await ctx.send("An unexpected error occurred. Please try again later.")

@bot.command(name='generate', aliases=['gen', 'img'])
async def generate_image(ctx, *, prompt: str):
    """Generate an image based on a text prompt using Gemini 2.0 Flash Preview."""
    
    # Check rate limiting
    if rate_limiter.is_rate_limited(ctx.author.id):
        reset_time = rate_limiter.get_reset_time(ctx.author.id)
        await ctx.send(f"Rate limit exceeded. Please try again in {reset_time} seconds.")
        return
    
    # Validate prompt length
    if len(prompt) > 1000:
        await ctx.send("Prompt is too long. Please keep it under 1000 characters.")
        return
    
    if len(prompt.strip()) < 3:
        await ctx.send("Prompt is too short. Please provide a more detailed description.")
        return
    
    # Send initial response
    async with ctx.typing():
        try:
            # Generate image
            logger.info(f"Generating image for user {ctx.author.id} with prompt: {prompt[:100]}...")
            image_data = await image_generator.generate_image_async(prompt)
            
            if not image_data:
                await ctx.send("Failed to generate image. The model might be experiencing issues. Please try again later.")
                return
            
            # Create Discord file object
            image_file = discord.File(io.BytesIO(image_data), filename="generated_image.png")
            
            # Create embed for better presentation
            embed = discord.Embed(
                title="Generated Image",
                description=f"**Prompt:** {prompt}",
                color=0x00ff00
            )
            embed.set_image(url="attachment://generated_image.png")
            embed.set_footer(text=f"Generated by {ctx.author.display_name}")
            
            # Send the image
            await ctx.send(embed=embed, file=image_file)
            logger.info(f"Successfully generated and sent image for user {ctx.author.id}")
            
        except Exception as e:
            logger.error(f"Error generating image: {e}")
            await ctx.send("Failed to generate image. Please check your prompt and try again later.")

@bot.command(name='help_generate', aliases=['genhelp'])
async def help_generate(ctx):
    """Show help information for the image generation command."""
    embed = discord.Embed(
        title="Image Generation Help",
        description="Generate images using Google's Gemini 2.0 Flash Preview model",
        color=0x0099ff
    )
    
    embed.add_field(
        name="Commands",
        value="`!generate <prompt>` - Generate an image\n`!gen <prompt>` - Short alias\n`!img <prompt>` - Another alias",
        inline=False
    )
    
    embed.add_field(
        name="Usage Tips",
        value="• Be descriptive in your prompts\n• Include style, mood, and details\n• Prompts should be 3-1000 characters\n• Use English for best results",
        inline=False
    )
    
    embed.add_field(
        name="Examples",
        value="`!generate a majestic dragon flying over a fantasy castle at sunset`\n`!gen cyberpunk cityscape with neon lights and rain`\n`!img cute cat wearing a wizard hat, digital art style`",
        inline=False
    )
    
    embed.add_field(
        name="Rate Limits",
        value=f"• {RATE_LIMIT_REQUESTS} requests per {RATE_LIMIT_WINDOW//60} minutes per user\n• Please be patient, image generation takes time",
        inline=False
    )
    
    await ctx.send(embed=embed)

@bot.command(name='ping')
async def ping(ctx):
    """Check bot latency."""
    latency = round(bot.latency * 1000)
    await ctx.send(f"Pong! Latency: {latency}ms")

@bot.command(name='info')
async def bot_info(ctx):
    """Show bot information."""
    embed = discord.Embed(
        title="Bot Information",
        color=0x0099ff
    )
    
    embed.add_field(name="Servers", value=len(bot.guilds), inline=True)
    embed.add_field(name="Users", value=len(set(bot.get_all_members())), inline=True)
    embed.add_field(name="Latency", value=f"{round(bot.latency * 1000)}ms", inline=True)
    
    embed.add_field(
        name="Features",
        value="• AI Image Generation\n• Rate Limiting\n• Error Handling\n• Multiple Command Aliases",
        inline=False
    )
    
    embed.set_footer(text="Powered by Google Gemini 2.0 Flash Preview")
    
    await ctx.send(embed=embed)

if __name__ == "__main__":
    try:
        bot.run(DISCORD_TOKEN)
    except discord.LoginFailure:
        logger.error("Invalid Discord token provided")
    except Exception as e:
        logger.error(f"Failed to start bot: {e}")