import os
from openai import OpenAI
import gradio as gr

# Initialize OpenAI client with Groq model
XAI_API_KEY = os.getenv("XAI_API_KEY")
client = OpenAI(
    api_key=XAI_API_KEY,
    base_url="https://api.x.ai/v1",
)

# Define a function to handle chat with humor, including chat history
def chat_with_buddy(history, user_message):
    # Start the conversation with a system message and add previous chat history
    messages = [{"role": "system", "content": "You are Chat Buddy, a friendly and humorous chatbot who loves to make people smile."}]
    messages += [{"role": "user", "content": msg[0]} for msg in history] + [{"role": "assistant", "content": msg[1]} for msg in history if len(msg) > 1]
    messages.append({"role": "user", "content": user_message})
    
    try:
        # Generate a response from Groq
        response =  client.chat.completions.create(
            model="grok-beta",
            messages=messages
        )
        # Extract the response content
        reply = response.choices[0].message.content


        # Update chat history with the new user message and assistant response
        history.append((user_message, reply))
        return history
    
    except Exception as e:
        # In case of error, add error message to history
        error_message = f"Oops, something went wrong! Just blame it on the robots... 😉 (Error: {str(e)})"
        history.append((user_message, error_message))
        return history

# Set up the Gradio chat interface with history
with gr.Blocks() as chat_interface:
    gr.Markdown("# Chat Buddy")
    gr.Markdown("Your friendly chatbot with a touch of humor! Ask Chat Buddy anything, and enjoy a chat that’s both helpful and light-hearted.")
    
    chatbot = gr.Chatbot()
    message = gr.Textbox(placeholder="Type your message here...", label="Your Message")
    submit_button = gr.Button("Send")

    # Define the submit action to update chat history
    def submit(user_message, chat_history):
        return chat_with_buddy(chat_history, user_message), ""

    # Link the submit button to the function
    message.submit(submit, inputs=[message, chatbot], outputs=[chatbot, message])
    submit_button.click(submit, inputs=[message, chatbot], outputs=[chatbot, message])

# Launch the chat interface
chat_interface.launch()