import discord
from discord import app_commands
from discord.ui import Button, View
import os
import io
import asyncio
import aiohttp
import gradio as gr

# --- Environment Variables & Setup ---
DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
GLIF_API_TOKEN = os.getenv("GLIF_API_TOKEN")
GLIF_API_URL = "https://simple-api.glif.app"

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

# --- Discord Bot Setup ---
intents = discord.Intents.default()
client = discord.Client(intents=intents)
tree = app_commands.CommandTree(client)

# --- Image Generation Function ---
async def generate_image_async(prompt, aspect_ratio):
    payload = {
        "id": "cm3ugmzv2002gnckiosrwk6xi",
        "inputs": [prompt, aspect_ratio],
    }
    headers = {"Authorization": f"Bearer {GLIF_API_TOKEN}"}

    async with aiohttp.ClientSession() as session:
        try:
            async with session.post(
                GLIF_API_URL, json=payload, headers=headers, timeout=15
            ) as response:
                response.raise_for_status()
                response_data = await response.json()
                if "output" in response_data:
                    return response_data["output"]
                elif "error" in response_data:
                    return f"Error: {response_data['error']}"
                else:
                    return "Error: Unexpected response from GLIF API."

        except asyncio.TimeoutError:
            return "Error: GLIF API request timed out."
        except aiohttp.ClientError as e:
            return f"API request failed: {e}"

# --- Discord Command: /generate ---
@tree.command(name="generate", description="Generates an image based on a text prompt")
@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"),
        app_commands.Choice(name="9:21", value="9:21"),
        app_commands.Choice(name="21:9", value="21:9"),
    ]
)
async def generate_command(
    interaction: discord.Interaction,
    prompt: str,
    aspect_ratio: app_commands.Choice[str],
):
    try:
        await interaction.response.defer()

        image_url_or_error = await generate_image_async(prompt, aspect_ratio.value)

        if image_url_or_error.startswith("http"):
            # Download the image
            async with aiohttp.ClientSession() as session:
                async with session.get(image_url_or_error) as resp:
                    if resp.status != 200:
                        return await interaction.followup.send('Could not download file...')
                    image_data = await resp.read()

            # Format the prompt for display with expand/collapse
            formatted_prompt, initial_content, full_content = format_prompt_for_discord(prompt)

            # Create buttons
            async def show_more_callback(interaction):
                await interaction.response.edit_message(content=full_content, view=view)

            async def show_less_callback(interaction):
                await interaction.response.edit_message(content=initial_content, view=view)

            show_more_button = Button(label="Read More", style=discord.ButtonStyle.primary)
            show_more_button.callback = show_more_callback
            show_less_button = Button(label="Hide", style=discord.ButtonStyle.secondary)
            show_less_button.callback = show_less_callback

            view = View()
            if len(prompt) > 256:
                view.add_item(show_more_button)
            view.add_item(show_less_button)

            # Send the image as an attachment
            with io.BytesIO(image_data) as image_file:
                await interaction.followup.send(
                    content=initial_content,
                    file=discord.File(image_file, 'generated_image.png'),
                    view=view
                )

        else:
            await interaction.followup.send(
                f"Sorry, I couldn't generate an image. {image_url_or_error}"
            )
    except Exception as e:
        await interaction.followup.send(f"Error: {str(e)}")

# --- Helper Function: Format Prompt for Discord ---
def format_prompt_for_discord(prompt):
    """Formats the prompt for display in Discord with expand/collapse functionality using buttons."""
    max_length = 256
    if len(prompt) <= max_length:
        return prompt, f"**Prompt:** {prompt}", f"**Prompt:** {prompt}"

    truncated_prompt = prompt[:max_length]
    initial_content = f"**Prompt:** {truncated_prompt}..."
    full_content = f"**Prompt:** {prompt}"

    return prompt, initial_content, full_content

# --- Discord Command: /hello ---
@tree.command(name="hello", description="Says hello!")
async def hello_command(interaction: discord.Interaction):
    await interaction.response.send_message("Hello there!")

# --- Bot Ready Event ---
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)

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