import streamlit as st
from llama_index.llms.gemini import Gemini
from llama_index.readers.web import BeautifulSoupWebReader
from llama_index.core import SummaryIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
import os
import google.generativeai as genai
from dotenv import load_dotenv
import uuid
from datetime import datetime, timedelta

# Page config with initial sidebar state
st.set_page_config(
    page_title="Web Content Query Assistant",
    layout="wide",
    initial_sidebar_state="expanded",  # This keeps the sidebar open by default
)

# Custom CSS with improved contrast
st.markdown(
    """
<style>
    /* Main container styling */
    .main {
        padding: 2rem;
    }

    /* Sidebar styling */
    .css-1d391kg {
        padding: 2rem 1rem;
    }

    /* Chat container styling */
    .stChatMessage {
        border-radius: 15px;
        padding: 1rem;
        margin: 0.5rem 0;
    }

    /* User message styling */
    .stChatMessage[data-testid="user-message"] {
        background-color: #1976D2 !important;
        color: white !important;
    }

    /* Assistant message styling */
    .stChatMessage[data-testid="assistant-message"] {
        background-color: #2F3336 !important;
        color: white !important;
    }

    /* Make all text in chat messages white */
    .stChatMessage p {
        color: white !important;
    }

    /* Input box styling */
    .stTextInput input {
        border-radius: 10px;
    }

    /* Button styling */
    .stButton button {
        border-radius: 20px;
        padding: 0.5rem 2rem;
        background-color: #2196F3;
        color: white;
        border: none;
        transition: all 0.3s ease;
    }

    .stButton button:hover {
        background-color: #1976D2;
        transform: translateY(-2px);
    }

    /* Error message styling */
    .stAlert {
        border-radius: 10px;
    }

    /* Chat input container styling */
    .stChatInputContainer {
        padding-top: 1rem;
        border-top: 1px solid #4a4a4a;
    }
</style>
""",
    unsafe_allow_html=True,
)

# Load environment variables
load_dotenv()


# Initialize the embedding model and LLM
@st.cache_resource
def init_models():
    embed_model = HuggingFaceEmbedding(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )
    llm = Gemini(model="models/gemini-2.0-flash-exp")
    return embed_model, llm


# Session state initialization
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "index_cache" not in st.session_state:
    st.session_state.index_cache = {}
if "current_url" not in st.session_state:
    st.session_state.current_url = ""


class IndexManager:
    def __init__(self, embed_model, llm):
        api_key = os.getenv("GOOGLE_API_KEY")
        if not api_key:
            st.error(
                "🔑 GOOGLE_API_KEY not found in environment variables. Please check your .env file."
            )
            st.stop()
        genai.configure(api_key=api_key)
        self.embed_model = embed_model
        self.llm = llm

    def process_url(self, url):
        if url in st.session_state.index_cache:
            return st.session_state.index_cache[url]

        with st.spinner("📚 Processing webpage content..."):
            try:
                documents = BeautifulSoupWebReader().load_data([url])
                index = SummaryIndex.from_documents(
                    documents, embed_model=self.embed_model
                )
                st.session_state.index_cache[url] = index
                return index
            except Exception as e:
                st.error(f"❌ Error processing URL: {str(e)}")
                return None

    def query_index(self, index, query):
        query_engine = index.as_query_engine(
            response_mode="tree_summarize", llm=self.llm
        )
        return query_engine.query(query)


def main():
    embed_model, llm = init_models()

    # Sidebar with enhanced styling
    with st.sidebar:
        st.markdown("# 🤖 Web Content Assistant")
        st.markdown("---")

        url = st.text_input(
            "🌐 Enter webpage URL:",
            placeholder="https://example.com",
            value=st.session_state.current_url,
        )

        if url != st.session_state.current_url:
            st.session_state.current_url = url
            st.session_state.chat_history = []

        st.markdown("---")

        # Settings section
        st.markdown("### ⚙️ Settings")
        if st.button("🗑️ Clear Chat History", use_container_width=True):
            st.session_state.chat_history = []
            st.rerun()

        # Help section
        st.markdown("### 📖 How to Use")
        st.markdown(
            """
        1. 🔗 Enter a webpage URL above
        2. 💭 Ask questions in the chat
        3. 🤖 Get AI-powered answers

        **Tips:**
        - Ask specific questions
        - One topic at a time
        - Clear chat to start fresh
        """
        )

    # Main chat container
    st.markdown("## 💬 Chat Interface")

    # Create a container for chat messages with fixed height and scrolling
    chat_container = st.container()
    with chat_container:
        for message in st.session_state.chat_history:
            with st.chat_message(
                message["role"], avatar="🧑‍💻" if message["role"] == "user" else "🤖"
            ):
                st.markdown(message["content"])

    # Chat input with error handling
    if prompt := st.chat_input("💭 Ask about the webpage content...", key="chat_input"):
        if not url:
            st.error("⚠️ Please enter a URL first!")
            st.stop()

        # Display user message
        with st.chat_message("user", avatar="🧑‍💻"):
            st.markdown(prompt)
        st.session_state.chat_history.append({"role": "user", "content": prompt})

        try:
            index_manager = IndexManager(embed_model, llm)
            index = index_manager.process_url(url)

            if index:
                with st.chat_message("assistant", avatar="🤖"):
                    with st.spinner("🤔 Thinking..."):
                        response = index_manager.query_index(index, prompt)
                        st.markdown(str(response))
                st.session_state.chat_history.append(
                    {"role": "assistant", "content": str(response)}
                )
            else:
                st.error(
                    "❌ Failed to process the webpage. Please check the URL and try again."
                )

        except Exception as e:
            st.error(f"❌ Error: {str(e)}")


if __name__ == "__main__":
    main()