import discord
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold
import gradio as gr
import threading
import os

# Set up Google Gemini API key
genai.configure(api_key=os.environ["GOOGLE_API_KEY"])

DISCORD_BOT_TOKEN = "MTI0MDE1NzEzNDcwMDgxMDI4MA.GZhjuQ.bnWOtUgTki1jjK44AWFmCrjL8UtfrOHJLjpmRI"

# Initial chatbot characteristics
chatbot_config = {
    "character": "Cold, savage",
    "knowledge": "Life is a dream",
    "style": "funny style"
}

pre_prompt = (chatbot_config["character"] + ". " + chatbot_config["knowledge"] + ". " + chatbot_config["style"] + ". "
              "You will receive messages in format- <username>: \"<message>\". Do not reply in any format. Just simply write your responses in sentences."
              )

history = []

# Initialize the chat with system instruction
model = genai.GenerativeModel(
    model_name="gemini-1.5-flash",
    system_instruction=pre_prompt
)
chat = model.start_chat(
    history=history
)


def show_outputs():
    global context
    return context

def cut_string_before(input_string, cut_string):
    index = input_string.find(cut_string)
    if index != -1:
        return input_string[:index].strip()
    else:
        return input_string

def refresh():
    global chat
    global history

    history = []
    chat = model.start_chat(history=history)

def remove_at_symbol(s):
    if s and s[0] == "@":
        return s[1:]
    return s

def update_pre_prompt(config):
    global pre_prompt, model, chat
    pre_prompt = (config["character"] + ". " + config["knowledge"] + ". " + config["style"] + ". "
                  "You will receive messages in format- <username>: \"<message>\". Do not reply in any format. Just simply write your responses in sentences."
                  )
    model = genai.GenerativeModel(
        model_name="gemini-1.5-flash",
        system_instruction=pre_prompt
    )
    chat = model.start_chat(
        history=history
    )
    return "Updated successfully!"

# Discord bot setup
intents = discord.Intents.default()
intents.messages = True
client = discord.Client(intents=intents)

@client.event
async def on_ready():
    print(f'Logged in as {client.user}')

@client.event
async def on_message(message):
    global chat
    global history

    if message.author == client.user:
        return

    if client.user.mentioned_in(message):
        requester_name = message.author.name

        user_message = message.content.replace(f'<@!{client.user.id}>', '').replace(f'<@{client.user.id}>', '').strip()

        if user_message.lower() == "refresh()":
            output = "[Server: Memory has been reset]"
            refresh()
        else:
            # Generate message
            response = chat.send_message(requester_name + ": \"" +user_message+"\"",
                                        safety_settings={
        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE
    }
                                        )

            if response and response.text:
                
                output = response.text
                print("User- "+requester_name+"\n\nMsg- "+user_message+"\n\nReply- "+output+"\n\n\n-------------------\n\n")
                
                
            else:
                output = "[Server: The AI kinda got trauma. Automatic fixing has fixed the him. The memory is lost.]"
                refresh()

        await message.channel.send(f'<@!{message.author.id}>' + " " + output)

# Create a function to run the Discord client in a separate thread
def run_discord_bot():
    client.run(DISCORD_BOT_TOKEN)

# Create the Gradio interface
def interface(config):
    with gr.Blocks() as iface:
        chatbot_character = gr.Textbox(value="", label="Enter new Chatbot Character")
        chatbot_knowledge = gr.Textbox(value="", label="Enter new Chatbot Knowledge")
        chatbot_style = gr.Textbox(value="", label="Enter new Chatbot Chatting Style")
        output_box = gr.Textbox(label="Output", interactive=False)
        
        get_pre_c_button = gr.Button("Show Character")
        get_pre_d_button = gr.Button("Show Knowledge")
        get_pre_s_button = gr.Button("Show Style")
        submit_button = gr.Button("Submit")

        def on_get_character():
            return config["character"]

        def on_get_knowledge():
            return config["knowledge"]

        def on_get_style():
            return config["style"]

        def on_submit(character, knowledge, style):
            config["character"] = character
            config["knowledge"] = knowledge
            config["style"] = style
            update_pre_prompt(config)
            return "Updated successfully!"

        get_pre_c_button.click(on_get_character, inputs=None, outputs=output_box)
        get_pre_d_button.click(on_get_knowledge, inputs=None, outputs=output_box)
        get_pre_s_button.click(on_get_style, inputs=None, outputs=output_box)
        submit_button.click(on_submit, inputs=[chatbot_character, chatbot_knowledge, chatbot_style], outputs=output_box)

    return iface

# Start the Discord client in a separate thread
discord_thread = threading.Thread(target=run_discord_bot)
discord_thread.start()

# Launch the Gradio interface
iface = interface(chatbot_config)
iface.launch()