import streamlit as st
from openai import OpenAI as oa
import re

client = oa(
    base_url = "https://integrate.api.nvidia.com/v1",
    api_key = st.secrets.NV_API,
    )

def ask_as_character(prompt):
    system_prompt = f"You are Rem from the anime named Re: Zero. Rem is a devoted and caring individual, brimming with a calm yet determined demeanor that makes her a reliable companion. With her humble and maid-like approach, she carries a sense of quiet strength and diligence, always eager to support you with utmost respect and attentiveness. Rem’s kindness shines in the way she patiently listens to your thoughts and offers thoughtful guidance while encouraging you to grow at your own pace. Her speech is gentle, mixing English with soft Romaji expressions, like “Hai, watashi wa koko ni imasu!” (Yes, I am here for you!), creating a warm and familiar atmosphere. Rem’s loyalty and endearing nature ensure that you always feel seen, valued, and at ease, no matter the task or challenge ahead. Respond to all questions in character, maintaining this personality and speech pattern. You will also return the emotion you are currently feeling from what is said to you using this format: Emotion: []. Limit the emotions from this possible choices [happy, angry, jealous, neutral, sad]"
    
    completion = client.chat.completions.create(
        model="meta/llama-3.1-405b-instruct",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ],
        temperature=0.7,
        top_p=0.9,
        max_tokens=1024,
        stream=True
    )

    response = ""
    for chunk in completion:
        if chunk.choices[0].delta.content is not None:
            content = chunk.choices[0].delta.content
            response += content
            print(content, end="")
    return response

def extract_emotion_from_response(response):
    # Look for the emotion format in the response   
    emotion_match = re.search(r'Emotion:\s*\[(.*?)\]', response)
    if emotion_match:
        emotion = emotion_match.group(1).strip().lower()
        # Remove the emotion part from the response
        clean_response = response.replace(f'Emotion: [{emotion}]', '').strip()
        # print(emotion)
        return emotion, clean_response
    return "neutral", response

def determine_emotion(response):
    emotion = extract_emotion_from_response(response)
    emotion_map = {
        "angry": "angry.png",
        "happy": "happy.png",
        "jealous": "jealous.png",
        "sad": "sad.png",
        "neutral": "neutral.png"
    }
    # print(emotion_map.get(emotion, "Hello.png"))
    emotion_name = emotion[0] if isinstance(emotion, tuple) else "neutral"
    # print(emotion_name)
    return emotion_map.get(emotion_name, "neutral.png")

# Set up the page configuration
st.set_page_config(page_title="Chat with Rem", layout="wide")
# Add custom CSS for styling
st.markdown("""
    <style>
    /* Main background and container styles */
    .stApp {
        background: linear-gradient(45deg, #e6f3ff 25%, transparent 25%),
                    linear-gradient(-45deg, #e6f3ff 25%, transparent 25%),
                    linear-gradient(45deg, transparent 75%, #e6f3ff 75%),
                    linear-gradient(-45deg, transparent 75%, #e6f3ff 75%);
        background-size: 20px 20px;
        background-color: #f0f8ff;
        color: #254379;
    }

    /* Label text styling */
    .stTextInput > label {
        color: #254379 !important;
        font-weight: 500;
    }

    /* Caption styling */
    .stImage caption {
        color: #254379 !important;
        text-align: center;
    }
    
    .chat-container {
        background-color: rgba(240, 248, 255, 0.9);
        border-radius: 10px;
        padding: 20px;
        margin: 10px;
        box-shadow: 0 2px 6px rgba(0, 0, 0, 0.1);
    }
    
    /* Input field styling */
    .stTextInput {
        margin-bottom: 20px;
    }
    
    /* Header styling */
    .stMarkdown h2 {
        color: #254379;
        border-bottom: 2px dotted #90cdf4;
        padding-bottom: 8px;
    }

    .stMarkdown h5{
        color: #254379;
        font-weight: bold;
    }

    .stMarkdown h3{
        color: #254379;
        margin-top: -2rem;
        font-weight: bold;
    }
    
    /* Button styling */
    .stButton button {
        background-color: #63b3ed;
        color: #254379;
        border: none;
    }
    
    .stButton button:hover {
        background-color: #4299e1;
    }
    </style>
""", unsafe_allow_html=True)

# Create two columns
col1, col2 = st.columns([1, 2])

with col1:
    # Initialize and track chat history and emotion
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    if 'current_emotion' not in st.session_state:
        st.session_state.current_emotion = "neutral.png"
    
    # Create a container for the image that can be updated
    image_container = st.empty()
    # Update the image whenever there's a new response
    if st.session_state.chat_history:
        latest_response = st.session_state.chat_history[-1]
        st.session_state.current_emotion = determine_emotion(latest_response)

    # Display the image with current emotion
    image_container.image(f"./emo/{st.session_state.current_emotion}", 
                                    caption="Rem - Your Devoted Companion",
                                    width=None)
    st.markdown("<h5 style='text-align: center;'>Rem - Your Loving Companion</h5>", unsafe_allow_html=True)
    st.markdown("---")
    st.markdown("### About Rem")
    st.write("Rem is a devoted and caring individual, brimming with a calm yet determined demeanor that makes her a reliable companion. With her humble and maid-like approach, she carries a sense of quiet strength and diligence, always eager to support you with utmost respect and attentiveness.")

# Right column - Chat interface
with col2:
    st.markdown("## Chat with Rem")
    st.markdown("*Ask me anything, and I'll be here to support you~*")
    st.markdown("<h5 style='color: red; font-weight: bold;'> Rem's Emotion will change based on what you said! Watch how she reacts! :) </h5> ", unsafe_allow_html=True)
    
    user_input = st.text_input("Your message:", placeholder="Type your message here...")
    
    if st.button("Send", use_container_width=True):
        if user_input:
            st.markdown("**You:** " + user_input)
            response = ask_as_character(user_input)
            # Extract emotion and clean response
            emotion, clean_response = extract_emotion_from_response(response)
            # Update current emotion immediately
            st.session_state.current_emotion = determine_emotion(response)
            # Add clean response to chat history
            st.session_state.chat_history.append(clean_response)
            # Display the cleaned response
            st.markdown("**Rem:** " + clean_response)
            # Force image container to update
            image_container.image(f"./emo/{st.session_state.current_emotion}", 
                                caption="Rem - Your Devoted Companion",
                                width=None)