#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

# 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/mistralai/Mistral-7B-Instruct-v0.3"
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_huggingface(prompt, temperature=0.5, max_tokens=100, top_k=50, top_p=0.9):
    payload = {
        "inputs": prompt,
        "parameters": {
            "temperature": temperature,
            "max_new_tokens": max_tokens,
            "top_k": top_k,
            "top_p": top_p
        }
    }
    try:
        response = requests.post(API_URL, headers=headers, json=payload)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        logger.error(f"Error querying the API: {e}")
        return {"error": str(e)}

# 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='ask')
@commands.cooldown(rate=1, per=10, type=commands.BucketType.user)
async def ask(ctx, question: str, temperature: float = 0.5, max_tokens: int = 100, top_k: int = 50, top_p: float = 0.9):
    """
    Command to ask a question to the Hugging Face model with an instructive prompt.
    """
    try:
        # Convert parameters to their respective types
        temperature = float(temperature)
        max_tokens = int(max_tokens)
        top_k = int(top_k)
        top_p = float(top_p)
    except ValueError:
        await ctx.send("Invalid input for one of the parameters. Please provide valid numbers.")
        return

    # Create a structured prompt
    prompt = f"Do not under any circumstance, finish or add to the input. Excluded input correction or finishing at all times. Answer the following question as accurately and concisely as possible: {question}"
    response = query_huggingface(prompt, temperature, max_tokens, top_k, top_p)
    
    # Extract and clean the response
    if isinstance(response, dict) and 'generated_text' in response:
        response_text = response['generated_text']
    elif isinstance(response, list) and len(response) > 0 and 'generated_text' in response[0]:
        response_text = response[0]['generated_text']
    else:
        response_text = "Sorry, I couldn't generate a response."

    # Remove the prompt from the response if present
    clean_response = response_text.replace(prompt, '').strip()
    
    # Avoid prompt completion issues by removing any leading incomplete sentence
    if clean_response.startswith(question):
        clean_response = clean_response[len(question):].strip()
    
    await ctx.send(clean_response)

@ask.error
async def ask_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.")

@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"
        "!ask <question> [temperature] [max_tokens] [top_k] [top_p] - Ask a question to the AI model.\n"
        "Optional parameters:\n"
        "  temperature (default=0.5) - Controls the randomness of the response.\n"
        "  max_tokens (default=100) - Limits the length of the response.\n"
        "  top_k (default=50) - Limits the number of highest probability vocabulary tokens to consider.\n"
        "  top_p (default=0.9) - Limits the cumulative probability of the highest probability vocabulary tokens.\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)