app.py

import discord
from discord import app_commands
import os
import requests
import asyncio
import aiohttp  # Use aiohttp for asynchronous HTTP requests
import gradio as gr  # Import Gradio
import google.generativeai as genai
from io import BytesIO
import PIL.Image
import base64
from typing import List
# --- Environment Variables & Setup ---
DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
# --- Discord Bot Setup ---
intents = discord.Intents.default()
client = discord.Client(intents=intents)
tree = app_commands.CommandTree(client)

genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel("gemini-2.0-flash-exp")

if not DISCORD_BOT_TOKEN or not GEMINI_API_KEY:
    raise ValueError("Both DISCORD_BOT_TOKEN and GEMINI_API_KEY must be set.")

@tree.command(name="hello", description="Says hello!")
async def hello_command(interaction):
    await interaction.response.send_message("Hello there!")

@tree.command(name="gemini", description="Chat with Gemini")
async def generate_command(
    interaction: discord.Interaction,
    prompt: str,
):
    try:
        await interaction.response.defer()  # Defer the interaction
        content = []  # List to store the content parts

        # Handle uploaded attachments
        if interaction.attachments:
            for attachment in interaction.attachments:
                # Download each attachment
                async with aiohttp.ClientSession() as session:
                    async with session.get(attachment.url) as resp:
                        if resp.status == 200:
                            image_bytes = await resp.read()
                            base64_image = base64.b64encode(image_bytes).decode("utf-8")
                            content.append({"mime_type": "image/jpeg", "data": base64_image})
                        else:
                            await interaction.followup.send(
                                f"Failed to download image: {attachment.filename}, status code: {resp.status}"
                            )
                            return

        # Add the text prompt to the content
        content.append(prompt)

        # Call the Generative AI model
        response = model.generate_content(content, stream=True)

        current_message = None
        current_message_content = ""

        # Process streaming responses
        for part in response:
            text_chunk = part.text

            if len(current_message_content) + len(text_chunk) <= 2000:
                current_message_content += text_chunk
                if current_message:
                    await current_message.edit(content=current_message_content)
                else:
                    current_message = await interaction.followup.send(content=current_message_content)
            else:
                if current_message:
                    await current_message.edit(content=current_message_content)
                else:
                    await interaction.followup.send(content=current_message_content)
                current_message_content = text_chunk
                current_message = await interaction.followup.send(content=current_message_content)

        if current_message_content:
            if current_message:
                await current_message.edit(content=current_message_content)
            else:
                await interaction.followup.send(content=current_message_content)

    except Exception as e:
        print(e)
        await interaction.followup.send(f"An error occurred: {e}")

@tree.command(name="prompt", description="Generate a prompt based on an image")
async def prompt_from_image(interaction: discord.Interaction):
    try:
        # Ensure the command is replying to a message
        if not interaction.message.reference:
            await interaction.response.send_message("Please reply to a message with an image to use this command.", ephemeral=True)
            return

        # Get the original message
        ref_message = await interaction.channel.fetch_message(interaction.message.reference.message_id)
        
        # Check if the message contains attachments
        if not ref_message.attachments:
            await interaction.response.send_message("The replied-to message does not contain any images.", ephemeral=True)
            return

        # Ensure the first attachment is an image
        attachment = ref_message.attachments[0]
        if not attachment.content_type or not attachment.content_type.startswith("image/"):
            await interaction.response.send_message("The replied-to message must contain an image.", ephemeral=True)
            return

        # Defer the interaction
        await interaction.response.defer()

        # Download the image
        async with aiohttp.ClientSession() as session:
            async with session.get(attachment.url) as resp:
                if resp.status == 200:
                    image_bytes = await resp.read()
                    base64_image = base64.b64encode(image_bytes).decode("utf-8")
                else:
                    await interaction.followup.send(
                        f"Failed to download the image: {attachment.filename}, status code: {resp.status}"
                    )
                    return

        # Call the Generative AI model to generate a description
        content = [{"mime_type": "image/jpeg", "data": base64_image}]
        
        # Ask Gemini to generate a description (prompt)
        response = model.generate_content(content, stream=False)

        # Collect the response text
        prompt = "".join([part.text for part in response])

        # Send the generated prompt back to the user
        await interaction.followup.send(f"Here is the suggested prompt for the image:\n```{prompt}```")

    except Exception as e:
        print(e)
        await interaction.followup.send(f"An error occurred: {e}")

async def on_ready():
    await tree.sync()
    print("Bot is ready!")


client.event(on_ready)

# --- Gradio Interface ---
def echo_text(text):
    return text


def run_gradio():
    gr.Interface(
        fn=echo_text,
        inputs="text",
        outputs="text",
        live=False,
        title="Minimal Gradio Interface",
    ).launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True)

async def main():
    bot_task = asyncio.create_task(client.start(DISCORD_BOT_TOKEN))
    gradio_task = asyncio.to_thread(run_gradio)

    await asyncio.gather(bot_task, gradio_task)


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