streamlit_app.py

"""
BOQTenders Streamlit Application

Interactive web interface for BOQ extraction and document chat.

Usage:
    streamlit run streamlit_app_new.py
"""
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent))

import tempfile
import streamlit as st
from loguru import logger

from config.settings import settings
from core.agent import BOQAgent

# Configure logging
logger.remove()
logger.add(
    sys.stderr,
    level=settings.log_level,
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
)


def initialize_services(api_key: str):
    """Initialize agent with API key (cached)."""
    if "agent_initialized" not in st.session_state or st.session_state.get("current_api_key") != api_key:
        st.session_state.agent = BOQAgent()
        st.session_state.agent_initialized = True
        st.session_state.current_api_key = api_key


def initialize_session_state():
    """Initialize Streamlit session state variables."""
    defaults = {
        "boq_output": None,
        "consistency": None,
        "qa_chain": None,
        "vector_store": None,
        "chunks": None,
        "chat_history": [],
        "document_loaded": False,
    }
    
    for key, value in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value


def process_pdf(uploaded_file, runs: int, boq_mode: list, specific_boq: str) -> bool:
    """
    Process uploaded PDF file using LangGraph agent.
    
    Args:
        uploaded_file: Streamlit uploaded file
        runs: Number of extraction runs
        boq_mode: List of BOQ modes ["default", "specific BOQ"]
        specific_boq: Specific BOQ string if applicable
    
    Returns:
        True if processing succeeded, False otherwise.
    """
    try:
        with st.spinner("Processing PDF..."):
            # Save to temp file
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:
                temp_file.write(uploaded_file.getvalue())
                temp_path = temp_file.name
            
            try:
                # Process with agent
                result = st.session_state.agent.process_document(
                    file_path=temp_path,
                    file_name=uploaded_file.name,
                    api_key=st.session_state.current_api_key,
                    runs=runs,
                    boq_mode=boq_mode,
                    specific_boq=specific_boq,
                    action="extract_boq"
                )
                
                if result["error"]:
                    st.error(f"Processing failed: {result['error']}")
                    return False
                
                # Store results in session
                st.session_state.boq_output = result["boq_output"]
                st.session_state.consistency = result["consistency"]
                st.session_state.qa_chain = result.get("qa_chain")  # Agent should return this
                st.session_state.document_loaded = True
                st.session_state.chat_history = []
                
                st.success("Document processed successfully")
                return True
                
            finally:
                Path(temp_path).unlink(missing_ok=True)
                
    except Exception as e:
        logger.error(f"Error processing PDF: {e}")
        st.error(f"Error processing PDF: {str(e)}")
        return False


def render_chat_interface():
    """Render the chat interface using agent."""
    st.subheader("💬 Chat with Document")
    
    if not st.session_state.document_loaded:
        st.info("Please upload a PDF to enable chat")
        return
    
    # Chat history
    for message in st.session_state.chat_history:
        role = message["role"]
        content = message["content"]
        
        if role == "user":
            st.chat_message("user").write(content)
        else:
            st.chat_message("assistant").write(content)
    
    # Chat input
    if prompt := st.chat_input("Ask a question about the document..."):
        st.session_state.chat_history.append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)
        
        with st.spinner("Thinking..."):
            try:
                # Use agent for chat
                answer = st.session_state.agent.chat_with_document(
                    process_id="streamlit_session",  # Use a fixed ID for Streamlit
                    question=prompt,
                    qa_chain=st.session_state.qa_chain,
                    chat_history=st.session_state.chat_history
                )
                
                st.session_state.chat_history.append({"role": "assistant", "content": answer})
                st.chat_message("assistant").write(answer)
                
            except Exception as e:
                logger.error(f"Chat error: {e}")
                error_msg = f"Error: {str(e)}"
                st.session_state.chat_history.append({"role": "assistant", "content": error_msg})
                st.chat_message("assistant").write(error_msg)


def render_boq_output():
    """Render the BOQ output."""
    st.subheader("📋 Extracted BOQ")
    
    if st.session_state.boq_output:
        st.markdown(st.session_state.boq_output)
        
        # Download button
        st.download_button(
            label="📥 Download BOQ as Markdown",
            data=st.session_state.boq_output,
            file_name="boq_output.md",
            mime="text/markdown"
        )
        
        # Consistency metrics
        if st.session_state.consistency:
            st.markdown("---")
            st.subheader("🔍 Consistency Metrics")
            
            consistency = st.session_state.consistency
            col1, col2, col3 = st.columns(3)
            
            with col1:
                st.metric("Consistency Score", f"{consistency['consistency_score']:.1f}%")
            
            with col2:
                st.metric("Avg Confidence", f"{consistency['avg_confidence']:.1f}%")
            
            with col3:
                st.metric("Successful Runs", f"{consistency['successful_runs']}/{consistency['runs']}")
            
            st.success("✅ Extraction completed successfully")
    else:
        st.info("Upload a PDF to see extracted BOQ items")


def render_consistency_check():
    """Render consistency check interface."""
    st.subheader("🔍 Consistency Check")
    
    if not st.session_state.document_loaded:
        st.info("Upload a PDF to run consistency checks")
        return
    
    runs = st.number_input(
        "Number of extraction runs",
        min_value=1,
        max_value=5,
        value=settings.consistency.default_runs,
        step=1,
        help="1 for quick check, 2-5 for comprehensive analysis"
    )
    
    if st.button("Run Consistency Check"):
        with st.spinner(f"Running {runs} extraction passes..."):
            try:
                result = st.session_state.consistency_checker.check(
                    chunks=st.session_state.chunks,
                    vector_store=st.session_state.vector_store,
                    runs=runs
                )
                
                col1, col2, col3 = st.columns(3)
                
                with col1:
                    st.metric("Consistency Score", f"{result['consistency_score']:.1f}%")
                
                with col2:
                    st.metric("Avg Confidence", f"{result['avg_confidence']:.1f}%")
                
                with col3:
                    st.metric("Successful Runs", f"{result['successful_runs']}/{result['runs']}")
                
                st.success("✅ Consistency check completed")
                    
            except Exception as e:
                logger.error(f"Consistency check error: {e}")
def render_sidebar():
    """Render the sidebar."""
    with st.sidebar:
        st.title("📄 BOQ Extractor")
        st.markdown("---")
        
        # API Key input
        api_key = st.text_input(
            "Google API Key",
            type="password",
            help="Enter your Google Generative AI API key",
            key="api_key_input"
        )
        
        if api_key:
            initialize_services(api_key)
        else:
            st.warning("Please enter your LLM API key to proceed.")
        
        st.markdown("---")
        
        # File upload
        uploaded_file = st.file_uploader(
            "Upload PDF Document",
            type=["pdf"],
            help="Upload a tender/BOQ document for extraction"
        )
        
        # Runs input
        runs_options = {
            1: "Quick (1 run) - Fast Execution",
            2: "Standard (2 runs) - Balanced Performance", 
            3: "Enhanced (3 runs) - Better Accuracy",
            4: "Precise (4 runs) - High Precision",
            5: "Maximum (5 runs) - Maximum Quality"
        }
        
        runs = st.selectbox(
            "Extraction Quality",
            options=list(runs_options.keys()),
            format_func=lambda x: runs_options[x],
            index=1,  # Default to 2 runs
            help="Choose extraction quality. Higher runs improve accuracy but take longer to process."
        )
        
        # BOQ Mode selection
        boq_mode = st.multiselect(
            "BOQ Extraction Mode",
            options=["default", "specific BOQ"],
            default=["default"],
            help="Select 'default' for all BOQ items, 'specific BOQ' to extract only a particular BOQ item."
        )
        
        specific_boq = None
        if "specific BOQ" in boq_mode:
            specific_boq = st.text_input(
                "Specific BOQ",
                help="Enter the name or description of the specific BOQ item to extract."
            )
        
        if uploaded_file and api_key:
            if st.button("🚀 Process Document"):
                process_pdf(uploaded_file, runs, boq_mode, specific_boq)
        elif uploaded_file and not api_key:
            st.error("Please enter API key first.")
        
        st.markdown("---")
        
        # Clear session
        if st.button("🗑️ Clear Session"):
            for key in list(st.session_state.keys()):
                if key not in ["agent_initialized", "current_api_key"]:
                    del st.session_state[key]
            initialize_session_state()
            st.success("Session cleared!")
            st.rerun()


def main():
    """Main application entry point."""
    # Page config
    st.set_page_config(
        page_title=settings.streamlit.page_title,
        page_icon=settings.streamlit.page_icon,
        layout=settings.streamlit.layout,
        initial_sidebar_state="expanded"
    )
    
    # Add CSS for sticky tabs
    st.markdown("""
        <style>
        /* Make tabs sticky at top */
        .stTabs [data-baseweb="tab-list"] {
            position: sticky;
            top: 0;
            background-color: white;
            z-index: 999;
            padding-top: 1rem;
            padding-bottom: 0.5rem;
            border-bottom: 1px solid #e6e6e6;
        }
        
        /* Dark mode support */
        @media (prefers-color-scheme: dark) {
            .stTabs [data-baseweb="tab-list"] {
                background-color: #0e1117;
                border-bottom: 1px solid #333;
            }
        }
        
        /* Streamlit dark theme */
        [data-theme="dark"] .stTabs [data-baseweb="tab-list"] {
            background-color: #0e1117;
            border-bottom: 1px solid #333;
        }
        </style>
    """, unsafe_allow_html=True)
    
    # Initialize
    initialize_session_state()
    
    # Render sidebar
    render_sidebar()
    
    # Main content tabs
    tab1, tab2 = st.tabs(["📋 BOQ Output", "💬 Chat"])
    
    with tab1:
        render_boq_output()
    
    with tab2:
        render_chat_interface()


if __name__ == "__main__":
    main()