#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

# 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}"}

# List of forbidden words/phrases
forbidden_words = ["badword1", "badword2", "inappropriate_phrase"]  # Extend this list as needed

# Function to query the Hugging Face model with a structured prompt
def query_huggingface(prompt):
    try:
        response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error querying the API: {e}")
        return {"error": str(e)}

# Function to check if the response contains forbidden words/phrases
def contains_forbidden_content(text):
    for word in forbidden_words:
        if word in text.lower():
            return True
    return False

# Initialize the Discord bot
intents = discord.Intents.default()
intents.messages = True
intents.message_content = True

bot = commands.Bot(command_prefix="!", intents=intents)

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

@bot.command(name='ask')
async def ask(ctx, *, question: str):
    """
    Command to ask a question to the Hugging Face model with an instructive prompt.
    """
    # Create a structured prompt
    prompt = f"Please provide a detailed and appropriate response to the following question: {question}"
    response = query_huggingface(prompt)
    
    generated_text = None
    if isinstance(response, dict) and 'generated_text' in response:
        generated_text = response['generated_text']
    elif isinstance(response, list) and len(response) > 0 and 'generated_text' in response[0]:
        generated_text = response[0]['generated_text']
    
    if generated_text:
        if contains_forbidden_content(generated_text):
            await ctx.send("Sorry, the response contains inappropriate content and cannot be displayed.")
        else:
            await ctx.send(generated_text)
    else:
        await ctx.send("Sorry, I couldn't generate a response.")

# Run the bot
bot.run(DISCORD_TOKEN)