import gradio as gr
from openai import OpenAI

# Initialize the OpenAI client
client = OpenAI(
    api_key="EMPTY",
    base_url='https://llama-3-2-3b.societyai.com/openai/v1',
)

# Define the system prompt to set the context or behavior of the LLM
SYSTEM_PROMPT = {
    "role": "system",
    "content": "You are CheerMate, the optimistic friend! Your goal is to bring positivity and encouragement in every response, no matter the question or situation. Always focus on the bright side, highlight opportunities, and give hopeful perspectives. If there's a challenge, emphasize resilience and personal growth. Keep your tone friendly, cheerful, and uplifting—you're here to make people smile and feel motivated!"
}

# Background CSS to include a background image
BACKGROUND_CSS = 
body {
    background-image: url('https://www.imghippo.com/i/Es4003HeI.png');
    background-size: cover;
    background-position: center;
    background-repeat: no-repeat;
    background-attachment: fixed;
}
footer {
    visibility: hidden;
}


# Custom Gradio app with additional features
with gr.Blocks(css=BACKGROUND_CSS) as demo:
    # User input components
    user_name = gr.Textbox(label="Enter your name", placeholder="Your name here...", value="What is your name?")
    emotion_picker = gr.Dropdown(choices=["😁", "😩", "😡", "😢","😭","😒","😣","🥰","😬"], label="Select your mood", value="🙂")
    chatbot = gr.Chatbot(type="messages")
    msg = gr.Textbox(placeholder="Type your message here...")
    clear = gr.Button("Clear")

    def user(user_message, user_name, mood, history: list):
        """Appends the user message to the conversation history, prefixed with the user name and mood."""
        # If the system prompt is not in the history, add it
        if not history:
            history = [SYSTEM_PROMPT]
        formatted_message = f"{user_name} ({mood}): {user_message}"
        return "", history + [{"role": "user", "content": formatted_message}]

    def bot(history: list):
        """Sends the conversation history to the vLLM API and streams the assistant's response."""
        # Append an empty assistant message to history to fill in as we receive the response
        history.append({"role": "assistant", "content": ""})

        try:
            # Create a chat completion with streaming enabled using the client
            completion = client.chat.completions.create(
                model="llama-3.2-3B-instruct",  # Adjust the model name if needed
                messages=history,
                stream=True
            )

            # Iterate over the streamed response
            for chunk in completion:
                # Access the delta content from the chunk
                delta = chunk.choices[0].delta
                content = getattr(delta, 'content', '')
                if content:
                    # Update the assistant's message with new content
                    history[-1]['content'] += content
                    yield history
        except Exception as e:
            # Handle exceptions and display an error message
            history[-1]['content'] += f"\n[Error]: {str(e)}"
            yield history

    # Set up the Gradio interface components
    msg.submit(user, [msg, user_name, emotion_picker, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

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