import os
import gradio as gr
from langchain_groq import ChatGroq
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationChain
from langchain.prompts import PromptTemplate

MODEL_NAME = "llama-3.3-70b-versatile"
DEFAULT_API_KEY = os.getenv("GROQ_API_KEY", "")

def initialize_chatbot(api_key, model_name=MODEL_NAME):
    llm = ChatGroq(
        groq_api_key=api_key,
        model_name=model_name,
        temperature=0.7,
        max_tokens=1024
    )
    
    memory = ConversationBufferMemory(
        return_messages=True,
        memory_key="history"
    )
    
    template = """You are a helpful AI assistant. Have a natural conversation with the user.

Current conversation:
{history}
Human: {input}
AI Assistant:"""
    
    prompt = PromptTemplate(
        input_variables=["history", "input"],
        template=template
    )
    
    return ConversationChain(
        llm=llm,
        memory=memory,
        prompt=prompt,
        verbose=False
    )

conversation_chain = None

def chat_function(message, api_key):
    global conversation_chain
    
    if not api_key:
        return "Please provide a Groq API key to start chatting."
    
    if conversation_chain is None:
        try:
            conversation_chain = initialize_chatbot(api_key)
        except Exception as e:
            return f"Error initializing chatbot: {str(e)}"
    
    try:
        return conversation_chain.predict(input=message)
    except Exception as e:
        return f"Error: {str(e)}"

def reset_conversation():
    global conversation_chain
    conversation_chain = None

with gr.Blocks(title="LLM based Chatbot") as demo:
    gr.Markdown("# 🤖 LLM based Chatbot")
    gr.Markdown("Chat with an AI assistant powered by LangChain and Groq")
    gr.Markdown(f"**Model:** `{MODEL_NAME}`")
    
    if not DEFAULT_API_KEY:
        api_key_input = gr.Textbox(
            label="Groq API Key",
            placeholder="Enter your Groq API key here...",
            type="password"
        )
    else:
        api_key_input = gr.Textbox(
            type="password",
            value=DEFAULT_API_KEY,
            visible=False
        )
    
    chatbot = gr.Chatbot(height=400)
    
    with gr.Row():
        msg = gr.Textbox(
            label="Message",
            placeholder="Type your message here...",
            scale=4
        )
        submit_btn = gr.Button("Send", scale=1)
    
    clear_btn = gr.Button("Clear Conversation")
    
    def respond(message, chat_history, api_key):
        if not message.strip():
            return chat_history, ""
        
        chat_history.append({"role": "user", "content": message})
        bot_message = chat_function(message, api_key)
        chat_history.append({"role": "assistant", "content": bot_message})
        
        return chat_history, ""
    
    def clear_chat():
        reset_conversation()
        return []
    
    msg.submit(respond, [msg, chatbot, api_key_input], [chatbot, msg])
    submit_btn.click(respond, [msg, chatbot, api_key_input], [chatbot, msg])
    clear_btn.click(clear_chat, None, chatbot)

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