import os
import time
from langchain_groq import ChatGroq
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from dotenv import load_dotenv
import gradio as gr

# Load environment variables
load_dotenv()
groq_api_key = os.getenv('GROQ_API_KEY')

if not groq_api_key:
    raise ValueError("GROQ_API_KEY not found in environment variables. Please add it to your .env file.")

# Global variables
stop_generation = False

# Initialize chatbot components
llm = ChatGroq(groq_api_key=groq_api_key, model_name="mixtral-8x7b-32768")
memory = ConversationBufferMemory()
conversation = ConversationChain(llm=llm, memory=memory)

def add_message(history, message):
    history.append((message, None))
    return history, gr.Textbox(value="", interactive=True)

def bot(history):
    global stop_generation
    stop_generation = False
    
    if not history:
        initial_message = "Hello! How can I assist you today?"
        history.append((None, initial_message))
        yield history
        return

    message = history[-1][0]
    
    try:
        response = conversation.predict(input=message)
    except Exception as e:
        response = f"I'm sorry, but I encountered an error: {str(e)}"
    
    history[-1] = (history[-1][0], "")
    for character in response:
        if stop_generation:
            break
        history[-1] = (history[-1][0], history[-1][1] + character)
        time.sleep(0.01)
        yield history

def stop_response():
    global stop_generation
    stop_generation = True

def clear_chat():
    global memory
    memory.clear()
    return []

with gr.Blocks() as demo:
    with gr.Column():
        chatbot = gr.Chatbot([], elem_id="chatbot")
        msg = gr.Textbox(label="Message")
        clear = gr.Button("Clear")

        chat_msg = msg.submit(add_message, [chatbot, msg], [chatbot, msg], queue=False).then(
            bot, chatbot, chatbot
        )

        clear.click(clear_chat, None, chatbot, queue=False)

        stop_btn = gr.Button("Stop Generation")
        stop_btn.click(stop_response)

    demo.queue()

if __name__ == "__main__":
    demo.launch()