import streamlit as st
from brave import perform_web_search, load_web_content, generate_detailed_explanation
import traceback
import sys

# Set page config with updated title
st.set_page_config(page_title="AI Search Engine 🔍", page_icon="🔍", layout="wide")

# Add this new CSS for the popup
st.markdown("""
<style>
    .stTextInput > div > div > input {
        font-size: 20px;
        border-radius: 25px;
        width: 100%;
    }
    .result-box {
        background-color: white;
        border: 1px solid #e0e0e0;
        border-radius: 10px;
        padding: 20px;
        margin-bottom: 20px;
    }
    .source-box {
        background-color: #f0f2f6;
        border-radius: 5px;
        padding: 10px;
        margin-bottom: 10px;
    }
    .thumbnail-image {
        width: 100%;
        height: 150px;
        object-fit: cover;
        border-radius: 10px;
        cursor: pointer;
    }
    .thumbnail-title {
        font-size: 12px;
        overflow-wrap: break-word;
        word-wrap: break-word;
        hyphens: auto;
        margin-top: 5px;
        margin-bottom: 15px;
        cursor: pointer;
    }
    .thumbnail-source {
        font-size: 10px;
        color: #888;
        margin-top: 5px;
    }
    .source-info {
        display: flex;
        align-items: center;
        font-size: 12px;
        color: #888;
        margin-top: 5px;
    }
    .source-info span {
        margin-right: 5px;
    }
    .query {
        font-size: 24px;
        font-weight: bold;
        margin-bottom: 10px;
    }
    .answer {
        font-size: 16px;
        margin-bottom: 20px;
    }
    .separator {
        border-top: 1px solid #e0e0e0;
        margin: 20px 0;
    }
    .popup-container {
        position: fixed;
        top: 0;
        left: 0;
        right: 0;
        bottom: 0;
        background-color: rgba(0, 0, 0, 0.5);
        display: flex;
        justify-content: center;
        align-items: center;
        z-index: 1000;
    }
    .popup-content {
        background-color: white;
        padding: 20px;
        border-radius: 10px;
        max-width: 500px;
        width: 100%;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'query_submitted' not in st.session_state:
    st.session_state.query_submitted = False
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'show_welcome' not in st.session_state:
    st.session_state.show_welcome = True

# Main content
st.title("AI Search Engine 🔍")
st.write("Ask any question and get detailed answers with sources.")

# Welcome popover
if st.session_state.show_welcome:
    with st.popover("Instructions"):
        st.markdown("""
        1. Only the first search performs a web search.
        2. The follow-up chat is for curating or modifying the response based on the initial search results.
        3. Follow-up questions do not perform additional web searches.
        4. To perform a new web search, click "Start New Search" at the bottom.
        """)
        st.markdown("Enjoy using AI Search Engine!")
        if st.button("Let's Search"):
            st.session_state.show_welcome = False
            st.rerun()

# Search input (only show if query hasn't been submitted)
if not st.session_state.query_submitted:
    query = st.text_input(
        label="Search Query",
        placeholder="Enter your question and press Enter",
        key="search_input",
        label_visibility="collapsed"
    )

    if query:
        st.session_state.query_submitted = True
        st.session_state.chat_history = []  # Reset chat history
        st.rerun()

# Main content area
if st.session_state.query_submitted:
    # Create two columns: one for the answer and chat, one for thumbnails
    col1, col2 = st.columns([3, 1])
    
    with col1:
        if not st.session_state.chat_history:
            # Initial query processing
            with st.spinner("Searching and analyzing..."):
                try:
                    search_results = perform_web_search(st.session_state.search_input)
                    web_content = load_web_content([result['url'] for result in search_results])
                    detailed_explanation = generate_detailed_explanation(st.session_state.search_input, web_content)
                    st.session_state.chat_history.append((st.session_state.search_input, detailed_explanation))
                    st.session_state.search_results = search_results
                except Exception as e:
                    st.error(f"An error occurred: {str(e)}")
                    st.error("Please try again or rephrase your question.")
                    st.session_state.query_submitted = False
                    st.rerun()

        # Display chat history
        for i, (q, a) in enumerate(st.session_state.chat_history):
            if i > 0:
                st.markdown('<div class="separator"></div>', unsafe_allow_html=True)
            st.markdown(f'<div class="query">{q}</div>', unsafe_allow_html=True)
            st.markdown(f'<div class="answer">{a}</div>', unsafe_allow_html=True)

        # Follow-up question input
        follow_up_key = f"follow_up_{len(st.session_state.chat_history)}"
        follow_up_query = st.text_input("Ask a follow-up question:", key=follow_up_key)

        if follow_up_query:
            with st.spinner("Generating response..."):
                try:
                    # Get the last 5 messages from the chat history
                    last_5_messages = st.session_state.chat_history[-5:]
                    context = "\n".join([f"Q: {q}\nA: {a}" for q, a in last_5_messages])
                    
                    follow_up_response = generate_detailed_explanation(follow_up_query, context)
                    st.session_state.chat_history.append((follow_up_query, follow_up_response))
                    st.rerun()  # Rerun to display the new response
                except Exception as e:
                    st.error(f"An error occurred while generating the follow-up response: {str(e)}")

        # Add a button to start a new search
        if st.button("Start New Search"):
            st.session_state.query_submitted = False
            st.session_state.chat_history = []
            st.rerun()

    with col2:
        # Display related images and sources
        st.markdown("### Related Images")
        for result in st.session_state.search_results:
            if result['thumbnail']:
                st.markdown(f"<a href='{result['url']}' target='_blank'><img src='{result['thumbnail']}' class='thumbnail-image'></a>", unsafe_allow_html=True)
                st.markdown(f"<a href='{result['url']}' target='_blank' style='text-decoration: none; color: inherit;'><div class='thumbnail-title'>{result['title']}</div></a>", unsafe_allow_html=True)
                st.markdown(f"""
                <div class='source-info'>
                    <span>Source: </span>
                    <a href='https://{result['hostname']}' target='_blank' style='text-decoration: none; color: #888;'>
                        {result['hostname']}
                    </a>
                </div>
                """, unsafe_allow_html=True)
            st.markdown("---")