import streamlit as st
from ai71 import AI71

# Access the API key from Streamlit secrets
ai71_api_key = st.secrets["AI71_API_KEY"]

# Initialize the AI71 client with the API key
client = AI71(ai71_api_key)

# Set page config with title and favicon
st.set_page_config(
    page_title="AI Lawyer 🧑‍⚖️",
    page_icon="assets/lawyer_icon.png",  # Replace with your favicon path
)

# Add custom CSS for styling
st.markdown(
    """
    <style>
    .main {
        background-color: #f0f4f8;
    }
    .sidebar .sidebar-content {
        background-color: #003366;
    }
    .stButton>button {
        color: #FFFFFF;
        background-color: #003366;
    }
    .stChatMessage--assistant {
        background-color: #e0f7fa;
    }
    .stChatMessage--user {
        background-color: #ffffff;
    }
    .title {
        color: #003366;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# Sidebar
st.sidebar.write("""
**AI Lawyer** is your intelligent assistant for legal advice and information. Powered by advanced AI technology, AI Lawyer helps you understand legal matters by providing detailed insights and potential solutions. Whether you have questions about contracts, disputes, or legal procedures, AI Lawyer is here to assist you.
""")

st.sidebar.header("How to Use AI Lawyer")
st.sidebar.write("""
1. **Enter Your Legal Question**:
   - Provide your legal query or describe the legal issue you are experiencing.

2. **Submit the Question**:
   - Use the input field at the bottom of the page to enter your query.

3. **Get a Response**:
   - AI Lawyer will process your input and generate a detailed response with relevant legal information and advice.

4. **Review and Take Action**:
   - Read the response provided by AI Lawyer and follow the suggested advice. Consult with a legal professional for further assistance if needed.
""")

# Show title and description.
st.markdown('<h1 class="title">AI Lawyer 🧑‍⚖️</h1>', unsafe_allow_html=True)
st.write(
    "This is your AI Lawyer that uses the AI71 model to provide legal advice and information."
)

# Initialize session state variables if not already set
if "messages" not in st.session_state:
    st.session_state.messages = []
    st.session_state.max_tokens = 512
    st.session_state.temperature = 0.7
    st.session_state.top_p = 0.95
    instruction = "Hi! This is your AI Lawyer 🧑‍⚖️. Please describe your legal question or issue. For example: 'I need help understanding a contract clause.'"
    st.session_state.messages.append({"role": "assistant", "content": instruction})

# Display the existing chat messages via st.chat_message.
for message in st.session_state.messages:
    if message["role"] == "assistant":
        with st.chat_message("assistant"):
            st.markdown(message["content"])
    elif message["role"] == "user":
        with st.chat_message("user"):
            st.markdown(message["content"])

# Create a chat input field to allow the user to enter a message. This will display
# automatically at the bottom of the page.
if prompt := st.chat_input("What legal question or issue do you need help with?"):

    # Store and display the current prompt.
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generate a response using the AI71 API.
    with st.spinner("Generating response..."):
        try:
            response = client.chat.completions.create(
                model="tiiuae/falcon-180B-chat",
                messages=[
                    {"role": "user", "content": prompt}
                ],
                max_tokens=st.session_state.max_tokens,
                temperature=st.session_state.temperature,
                top_p=st.session_state.top_p
            )
            # Collect and concatenate response chunks
            if response.choices and response.choices[0].message:
                full_response = response.choices[0].message.content

                # Stream the full response to the chat using st.write
                with st.chat_message("assistant"):
                    st.markdown(full_response)
                
                st.session_state.messages.append({"role": "assistant", "content": full_response})
        except Exception as e:
            st.error(f"An error occurred: {e}")