import streamlit as st
from rag_app import WebRAG
import time
import os

# Page Configuration
st.set_page_config(
    page_title="Web RAG Assistant",
    page_icon="🌐",
    layout="wide"
)

st.markdown(
    """
    <style>
        .stApp {
            max-width: 1200px;
            margin: 0 auto;
            font-family: 'Arial', sans-serif;
        }
        .sidebar .sidebar-content {
            background-color: #1E1E1E;
            color: white;
            padding: 20px;
            position: fixed;
            left: 0;
            top: 0;
            height: 100vh;
            width: 280px;
            z-index: 999;
        }
        .css-1d391kg {
            padding: 0;
        }
        .main-content {
            margin-left: 300px;
            padding: 20px;
        }
        .stTextInput, .stSelectbox, .stButton, .stChatInput {
            border-radius: 10px;
            padding: 10px;
        }
        .chat-container {
            border-radius: 10px;
            padding: 20px;
            background-color: #f0f2f6;
            margin: 10px 0;
        }
        .user-message {
            background-color: #2e7bf3;
            color: white;
            padding: 15px;
            border-radius: 15px;
            margin: 5px 0;
            text-align: right;
        }
        .assistant-message {
            background-color: white;
            padding: 15px;
            border-radius: 15px;
            margin: 5px 0;
            border: 1px solid #e0e0e0;
            text-align: left;
        }
    </style>
    """,
    unsafe_allow_html=True
)

# Initialize session state
if 'rag' not in st.session_state:
    st.session_state.rag = WebRAG()
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'url_processed' not in st.session_state:
    st.session_state.url_processed = False
if 'current_url' not in st.session_state:
    st.session_state.current_url = ""

# Function to reset chat history
def reset_chat_history():
    st.session_state.chat_history = []
    st.session_state.current_url = ""

# Sidebar on the extreme left
with st.sidebar:
    st.header("⚙️ Settings")
    url = st.text_input("Enter webpage URL:")
    
    scraping_method = st.selectbox(
        "Select Scraping Method",
        ["beautifulsoup", "scrapegraph", "crawl4ai"],
        help="""
        - BeautifulSoup: Basic HTML parsing, faster but less sophisticated
        - ScrapeGraph: AI-powered scraping, better at understanding content but slower
        - Crawl4ai: Advanced async crawler with good JavaScript support
        """
    )
    
    if st.button("🚀 Process URL"):
        if url:
            if url != st.session_state.current_url:
                reset_chat_history()
                st.session_state.current_url = url
            
            with st.spinner("Processing URL... Please wait."):
                try:
                    st.session_state.rag.crawl_and_process(url, scraping_method)
                    st.session_state.url_processed = True
                    st.success("✅ URL processed successfully!")
                    st.rerun()
                except Exception as e:
                    st.error(f"❌ Error processing URL: {str(e)}")
        else:
            st.warning("⚠️ Please enter a URL")
    
    st.divider()
    st.markdown("### 📌 How to use")
    st.markdown(
        """
        1. Enter a webpage URL
        2. Click 'Process URL' to analyze the content
        3. Ask questions about the webpage
        4. Receive AI-powered answers
        """
    )

# Main Content
st.markdown('<div class="main-content">', unsafe_allow_html=True)
st.title("🌐 Web RAG Assistant")
st.markdown("### Ask questions about any webpage")
st.divider()

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

# Chat Input
if st.session_state.url_processed:
    question = st.chat_input("💬 Ask a question about the webpage...")
    if question:
        st.session_state.chat_history.append({"role": "user", "content": question})
        with st.spinner("Thinking..."):
            try:
                answer = st.session_state.rag.ask_question(
                    question,
                    [(msg["content"], msg["content"]) for msg in st.session_state.chat_history if msg["role"] == "assistant"]
                )
                st.session_state.chat_history.append({"role": "assistant", "content": answer})
                st.rerun()
            except Exception as e:
                st.error(f"Error: {str(e)}")
else:
    st.info("👈 Please process a URL first using the sidebar")
st.markdown('</div>', unsafe_allow_html=True)