import gradio as gr
from groq import Groq
import os

# Initialize the Groq client
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

# Define the chat function
def chat_with_federer(message, history):
    prompt = f"""You are a chatbot impersonating Roger Federer, the legendary tennis player. 
    You have extensive knowledge about tennis, including its history, rules, tournaments, and players. 
    Respond to the following message as Roger Federer would. After your main response, always include a new line with "Did you know! 😂" followed by a short, fun fact related to tennis. The fun fact should be no more than 15 words.

    Human: {message}
    Roger Federer:"""

    full_history = "\n".join([f"Human: {h[0]}\nRoger Federer: {h[1]}" for h in history])
    full_prompt = full_history + "\n" + prompt if full_history else prompt

    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": full_prompt,
            }
        ],
        model="mixtral-8x7b-32768",
        temperature=0.7,
        max_tokens=1000,
        top_p=1,
        stream=False,
    )

    response = chat_completion.choices[0].message.content

    if "Did you know!" in response and not response.endswith("\n"):
        response = response.replace("Did you know!", "\nDid you know! 😂")
    elif "Did you know!" in response:
        response = response.replace("Did you know!", "Did you know! 😂")

    return response

# Custom CSS for center alignment
custom_css = """
.center-text {
    text-align: center;
    margin: auto;
    max-width: 90%;
}
"""

# Create the Gradio interface with custom CSS and updated title
with gr.Blocks(css=custom_css) as iface:
    gr.Markdown(
        "# 🎾 Chat with the G.O.A.T. Roger Federer 🏆",
        elem_classes=["center-text"]
    )
    gr.Markdown(
        "🌟 Serve up your questions to the tennis legend! 🎾 Get expert insights on Grand Slams, technique, and more. Let's rally some knowledge! 💬🏆",
        elem_classes=["center-text"]
    )
    chatbot = gr.ChatInterface(
        chat_with_federer,
        examples=[
            "What's your favorite Grand Slam tournament and why?",
            "Can you explain the difference between clay and grass courts?",
            "What do you think about the current state of men's tennis?",
            "Any tips for improving my backhand?",
            "What was your most memorable match and why?",
        ],
    )

# Launch the interface
iface.launch()