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
# --- 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)

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="imagen3", description="Generates iamge(s) using imagen3 model")
@app_commands.choices(
    aspect_ratio=[
        app_commands.Choice(name="1:1 (Square)", value="1:1"),
        app_commands.Choice(name="9:16 (Vertical)", value="9:16"),
        app_commands.Choice(name="16:9 (Horizontal)", value="16:9"),
        app_commands.Choice(name="3:4", value="3:4"),
        app_commands.Choice(name="4:3", value="4:3"),
    ]
)
async def generate_command(
    interaction: discord.Interaction,
    prompt: str,
    img_count:int,
    aspect_ratio: app_commands.Choice[str],
):
    genai.configure(api_key=GEMINI_API_KEY)
    imagen = genai.ImageGenerationModel("imagen-3.0-generate-001")

    result = imagen.generate_images(
        prompt=prompt,
        number_of_images=img_count,
        safety_filter_level="block_only_high",
        person_generation="allow_adult",
        aspect_ratio=aspect_ratio,
    )
    files = []
    for i, image in enumerate(result.generated_images):
        # Assuming `image` is a byte array or binary stream of the image
        image_file = BytesIO(image)  # Convert to BytesIO
        files.append(discord.File(fp=image_file, filename=f"image_{i+1}.png"))

    # Send images to Discord
    await interaction.followup.send(content=f"Here are your {img_count} images for prompt: `{prompt}`", files=files)
    
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())