import gradio as gr
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_groq import ChatGroq

# Session-wide context
conversation_history = [SystemMessage(content="You are a comedian AI assistant")]

def chatbot_response(api_key, user_input):
    if not api_key:
        return "Please enter your Groq API key.", ""

    # Pass API key to ChatGroq
    chat = ChatGroq(model="gemma2-9b-it", temperature=0.5, api_key=api_key)

    # Append user input and get response
    conversation_history.append(HumanMessage(content=user_input))
    response = chat.invoke(conversation_history)
    conversation_history.append(AIMessage(content=response.content))

    # Format chat history
    chat_log = ""
    for msg in conversation_history:
        if isinstance(msg, SystemMessage):
            chat_log += f"⚙️ *System Role*: {msg.content}\n\n"
        elif isinstance(msg, HumanMessage):
            chat_log += f"👤 **User**: {msg.content}\n\n"
        elif isinstance(msg, AIMessage):
            chat_log += f"🤖 **Bot**: {msg.content}\n\n"

    return response.content, chat_log

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Conversational Q&A Chatbot")
    # gr.Markdown("Enter your Groq API key and chat below!")
    gr.Markdown(
    "Enter your Groq API key and chat below! 🤖 This chatbot has been initialized as a **comedian AI assistant** — so expect jokes, quirky answers, and playful tone by default. You can repurpose it easily for serious Q&A, guidance, or custom personas!"
)

    api_key_box = gr.Textbox(label="🔑 Groq API Key", type="password")
    user_input = gr.Textbox(label="💬 Your Question")
    submit_btn = gr.Button("Ask")

    response_output = gr.Textbox(label="🧠 Response", interactive=False)
    chat_history = gr.Markdown()

    submit_btn.click(
        fn=chatbot_response,
        inputs=[api_key_box, user_input],
        outputs=[response_output, chat_history]
    )

demo.launch()