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):
    # Construct the prompt
    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 a short, funny "Did you know!" anecdote related to tennis. The anecdote should be no more than 15 words.

    Human: {message}
    Roger Federer:"""

    # Get the full conversation history
    full_history = "\n".join([f"Human: {h[0]}\nRoger Federer: {h[1]}" for h in history])
    
    # Combine history with the new prompt
    full_prompt = full_history + "\n" + prompt if full_history else prompt

    # Call the Groq API
    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,
    )

    # Extract the response
    response = chat_completion.choices[0].message.content

    # Ensure the "Did you know!" part is on a new line
    if "Did you know!" in response and not response.endswith("\n"):
        response = response.replace("Did you know!", "\nDid you know!")

    return response

# Create the Gradio interface
iface = gr.ChatInterface(
    chat_with_federer,
    title="🎾 Chat with the G.O.A.T. Roger Federer 🏆",
    description="🌟 Serve up your questions to the tennis legend! 🎾 Get expert insights on Grand Slams, technique, and more. Let's rally some knowledge! 💬🏆",
    theme="soft",
    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()