#import gradio as gr
#gr.load("models/mistralai/Mistral-7B-Instruct-v0.3").launch()
import os
import requests
import discord
from discord.ext import commands
from dotenv import load_dotenv
import logging
import io
import time
from PIL import Image

# Load environment variables from the .env file
load_dotenv()

DISCORD_TOKEN = os.getenv('dsTOK')
HF_API_KEY = os.getenv('HFREAD')

API_URL = "https://api-inference.huggingface.co/models/Corcelio/mobius"
headers = {"Authorization": f"Bearer {HF_API_KEY}"}

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Function to query the Hugging Face model with a structured prompt
def query(payload, max_retries=5, delay=5):
    retries = 0
    while retries < max_retries:
        try:
            response = requests.post(API_URL, headers=headers, json=payload)
            response.raise_for_status()
            return response.content
        except requests.exceptions.RequestException as e:
            logger.error(f"Error querying the API: {e}. Retrying in {delay} seconds...")
            time.sleep(delay)
            retries += 1
    logger.error(f"Failed to query the API after {max_retries} attempts.")
    return None

# Initialize the Discord bot with a default prefix
intents = discord.Intents.default()
intents.messages = True
intents.message_content = True
default_prefix = "!"
bot = commands.Bot(command_prefix=commands.when_mentioned_or(default_prefix), intents=intents)

@bot.event
async def on_ready():
    logger.info(f'Bot is ready. Logged in as {bot.user}')

@bot.command(name='generate')
@commands.cooldown(rate=1, per=10, type=commands.BucketType.user)
async def generate(ctx, *description):
    """
    Command to generate an image based on a description using the Hugging Face model.
    """
    if not description:
        await ctx.send("You need to provide a description for the image.")
        return

    description = " ".join(description)
    try:
        payload = {"inputs": description}
        image_bytes = query(payload)

        if image_bytes:
            image = Image.open(io.BytesIO(image_bytes))
            with io.BytesIO() as image_binary:
                image.save(image_binary, 'PNG')
                image_binary.seek(0)
                await ctx.send(file=discord.File(fp=image_binary, filename='generated_image.png'))
        else:
            await ctx.send("An error occurred while generating the image. Please try again later.")
    except Exception as e:
        logger.error(f"Error processing the generate command: {e}")
        await ctx.send("An error occurred while processing your request. Please try again later.")

@generate.error
async def generate_error(ctx, error):
    if isinstance(error, commands.CommandOnCooldown):
        await ctx.send(f"This command is on cooldown. Please try again after {int(error.retry_after)} seconds.")
    else:
        logger.error(f"Error in generate command: {error}")
        await ctx.send("An error occurred. Please try again later.")

@bot.command(name='setprefix')
@commands.has_permissions(administrator=True)
async def set_prefix(ctx, prefix: str):
    bot.command_prefix = commands.when_mentioned_or(prefix)
    await ctx.send(f"Command prefix changed to: {prefix}")

@bot.command(name='help_custom')
async def help_custom(ctx):
    help_text = (
        "Here are the commands you can use:\n"
        "!generate <description> - Generate an image based on the description.\n"
        "!setprefix <prefix> - Change the command prefix (admin only).\n"
        "!help_custom - Display this help message."
    )
    await ctx.send(help_text)

# Run the bot
bot.run(DISCORD_TOKEN)