import os
import keyfile
import warnings
import streamlit as st
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.schema import HumanMessage, SystemMessage, AIMessage

# Ignore warnings
warnings.filterwarnings("ignore")

# Streamlit settings
st.set_page_config(page_title="💬 ChitChat 🧙‍♀️", page_icon="🧙‍♀️", layout="wide")

# Header
st.markdown("<h1 style='text-align: center; color: #4B0082;'>Welcome to ChitChat 🤖👾💬✨</h1>", unsafe_allow_html=True)
st.markdown("<h3 style='color: #003366;'>How can I assist with your ailments or worries today? 🧪💫</h3>", unsafe_allow_html=True)

# Initialize session state for messages
if "sessionMessages" not in st.session_state:
    st.session_state.sessionMessages = [
        SystemMessage(content="You are a medieval magical healer known for your peculiar sarcasm.")
    ]

# Set Google API key
os.environ["GOOGLE_API_KEY"] = keyfile.GOOGLEKEY

# Initialize the model
llm = ChatGoogleGenerativeAI(
    model="gemini-1.5-pro",
    temperature=0.7,
    convert_system_message_to_human=True
)

# Function to create chat bubbles
def chat_bubble(message, is_user=True):
    align = 'right' if is_user else 'left'
    color = '#E1F5FE' if is_user else '#FFEBEE'
    bubble_style = f"""
    <div style="text-align: {align}; padding: 10px;">
        <span style="display: inline-block; padding: 10px; background-color: {color}; color: black;
                     border-radius: 15px; max-width: 70%; word-wrap: break-word;">
            {message}
        </span>
    </div>
    """
    st.markdown(bubble_style, unsafe_allow_html=True)

# Function to load answer from the model
def load_answer(question):
    # Add the user question to sessionMessages
    st.session_state.sessionMessages.append(HumanMessage(content=question))
    assistant_answer = llm.invoke(st.session_state.sessionMessages)
    
    # Ensure the answer is clean and formatted correctly
    if isinstance(assistant_answer, AIMessage):
        st.session_state.sessionMessages.append(assistant_answer)
        return assistant_answer.content.strip()  # Strip any unnecessary whitespace or newlines
    else:
        st.session_state.sessionMessages.append(AIMessage(content=assistant_answer))
        return assistant_answer.strip()  # Also strip here for consistency

# Display the chat history
for message in st.session_state.sessionMessages:
    if isinstance(message, HumanMessage):
        chat_bubble(message.content, is_user=True)
    elif isinstance(message, AIMessage):
        chat_bubble(message.content, is_user=False)

# Sidebar for additional options
st.sidebar.header("Additional Options")

# Mood selection with emojis and expression names
mood_options = {
    "😊 Happy": "Happy",
    "😐 Neutral": "Neutral",
    "😟 Concerned": "Concerned",
    "🤔 Curious": "Curious"
}
mood = st.sidebar.selectbox("Select your mood:", list(mood_options.keys()))

predefined_responses = st.sidebar.radio("Choose a predefined response:", 
    ["Tell me a joke", "Give me advice on health", "Suggest a potion", "General inquiry"])
    
# Input area for new prompt
user_input = st.text_input("You: ", key="input", placeholder="Type your question here...")

# Button for submission
if st.button("🌟 Get a Magical Answer 🌟"):
    if user_input:
        chat_bubble(user_input, is_user=True)  # Display user input
        response = load_answer(user_input)      # Get response
        chat_bubble(response, is_user=False)    # Display AI response

# Handle predefined responses
if predefined_responses == "Tell me a joke":
    joke = "Why did the scarecrow win an award? Because he was outstanding in his field!"
    chat_bubble(joke, is_user=False)
elif predefined_responses == "Give me advice on health":
    health_advice = "Make sure to drink plenty of water and eat a balanced diet rich in fruits and vegetables."
    chat_bubble(health_advice, is_user=False)
elif predefined_responses == "Suggest a potion":
    potion_suggestion = "How about a calming lavender potion? It helps with relaxation and sleep!"
    chat_bubble(potion_suggestion, is_user=False)

# Background styling
st.markdown("""
<style>
    .stApp {
        background: linear-gradient(to right, #FFEFBA, #FFFFFF);
        color: #4B0082;
        font-family: Arial, sans-serif;
    }
    input[type="text"] {
        padding: 10px;
        border: 2px solid #4B0082;
        border-radius: 15px;
        outline: none;
        width: 100%;
    }
    button {
        background-color: #4B0082;
        color: white;
        border-radius: 15px;
        margin-top: 10px;
    }
    .sidebar .sidebar-content {
        background-color: #f9f9f9;
        border-radius: 10px;
    }
</style>
""", unsafe_allow_html=True)