import streamlit as st
import google.generativeai as genai
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

# Configure Google Generative AI with API key
api_key = os.getenv("GENERATIVEAI_API_KEY")
genai.configure(api_key=api_key)

# Initialize session state for message history if it doesn't exist
if 'messages' not in st.session_state:
    st.session_state.messages = []

# Initialize session state for chat
if 'chat' not in st.session_state:
    st.session_state.chat = None

# Generation configuration and safety settings
generation_config = {
    "temperature": 0.9,
    "top_p": 0.5,
    "top_k": 5,
    "max_output_tokens": 1000,
}

safety_settings = [
    {
        "category": "HARM_CATEGORY_HARASSMENT",
        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
    },
    {
        "category": "HARM_CATEGORY_HATE_SPEECH",
        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
    },
    {
        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
    },
    {
        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
    },
]

def text_summary(text, isNew=False):
    if isNew or st.session_state.chat is None:
        model = genai.GenerativeModel(
            model_name="gemini-pro",
            generation_config=generation_config,
            safety_settings=safety_settings
        )
        st.session_state.chat = model.start_chat()
        st.session_state.chat.send_message("""
        Act as a financial advisor and generate financial summaries in a structured and tabular format...
        """)  # Your existing prompt here

    response = st.session_state.chat.send_message(text)
    return response.text

# Custom CSS for chat interface
st.markdown("""
<style>
.chat-container {
    height: 600px;
    overflow-y: auto;
    padding: 20px;
    margin-bottom: 100px;
}
.user-message {
    background-color: #e6f3ff;
    padding: 10px;
    border-radius: 10px;
    margin: 10px 0;
    max-width: 80%;
    margin-left: auto;
}
.bot-message {
    background-color: #f0f0f0;
    padding: 10px;
    border-radius: 10px;
    margin: 10px 0;
    max-width: 80%;
}
.input-container {
    position: fixed;
    bottom: 0;
    left: 0;
    right: 0;
    padding: 20px;
    background-color: white;
    border-top: 1px solid #ddd;
}
.stTextInput input {
    border-radius: 20px;
}
</style>
""", unsafe_allow_html=True)

# Main title
st.title("Financial Summary Chatbot")

# Chat message container
chat_container = st.container()

# Function to process the message
def process_message(message: str):
    if message.strip():
        # Add user message to history
        st.session_state.messages.append({"role": "user", "content": message})
        
        # Get bot response
        response = text_summary(message)
        
        # Add bot response to history
        st.session_state.messages.append({"role": "assistant", "content": response})

# Create a container for the input area at the bottom
with st.container():
    st.markdown('<div class="input-container">', unsafe_allow_html=True)
    
    # Create two columns for input and button
    col1, col2 = st.columns([5,1])
    
    with col1:
        user_input = st.text_input("Message", key="user_input", label_visibility="collapsed")
    
    with col2:
        if st.button("Send"):
            process_message(user_input)
            # Instead of trying to clear the input directly, we'll use a rerun
            st.rerun()
    
    st.markdown('</div>', unsafe_allow_html=True)

# Handle Enter key press
if user_input and len(user_input.strip()) > 0:
    if '\n' in user_input or st.session_state.get('enter_pressed', False):
        process_message(user_input)
        st.session_state.enter_pressed = False
        st.rerun()

# Display chat messages in the container
with chat_container:
    st.markdown('<div class="chat-container">', unsafe_allow_html=True)
    for message in st.session_state.messages:
        if message["role"] == "user":
            st.markdown(f'<div class="user-message">{message["content"]}</div>', unsafe_allow_html=True)
        else:
            st.markdown(f'<div class="bot-message">{message["content"]}</div>', unsafe_allow_html=True)
    st.markdown('</div>', unsafe_allow_html=True)