import streamlit as st
import logging
import asyncio
from contextlib import asynccontextmanager
from app import QueryRequest  # Import the request model

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Set page config
st.set_page_config(page_title="Certification Chat", layout="wide")
st.title("🎓 Hydrogen Certification Chat Assistant")

# Create a function to handle the async call
async def async_query(query_text):
    from app import handle_query  # Import here to avoid circular imports
    request = QueryRequest(query=query_text)
    return await handle_query(request)

# Function to run async code in Streamlit
def run_async(coroutine):
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    return loop.run_until_complete(coroutine)

# User input
user_input = st.text_area("💬 Enter your question about hydrogen certification:", height=100)

if user_input:
    try:
        # Use try-except to handle errors
        with st.spinner("Processing your query..."):
            # Run the async function
            result = run_async(async_query(user_input))
            
            # Display output in tabs
            st.markdown("## 🧠 Response")
            
            # Certification info
            st.subheader(f"Detected Certification: {result['certification']}")
            
            # Create tabs for answers and contexts
            tab1, tab2, tab3, tab4 = st.tabs([
                "Answers", 
                "Context Details", 
                "Raw Context (Dot Chunking)", 
                "Raw Context (Hybrid Chunking)"
            ])
            
            with tab1:
                st.markdown("### Basic Chunking Answer:")
                st.write(result["certif_index"])
                
                st.markdown("### Hybrid Chunking Answer:")
                st.write(result["certification_index"])
            
            with tab2:
                col1, col2 = st.columns(2)
                
                with col1:
                    st.markdown("### Basic Chunking Context Sources:")
                    for i, context_item in enumerate(result["context_certif"]):
                        with st.expander(f"Source {i+1}"):
                            st.write(context_item)
                
                with col2:
                    st.markdown("### Hybrid Chunking Context Sources:")
                    for i, context_item in enumerate(result["context_certifications"]):
                        with st.expander(f"Source {i+1}"):
                            st.write(context_item)
            
            with tab3:
                st.markdown("### Raw Context (Dot Chunking):")
                for i, chunk in enumerate(result["context_certif"]):
                    st.text_area(f"Chunk {i+1}", chunk, height=150)
            
            with tab4:
                st.markdown("### Raw Context (Hybrid Chunking):")
                for i, chunk in enumerate(result["context_certifications"]):
                    st.text_area(f"Chunk {i+1}", chunk, height=150)
                    
            # Add a section for feedback
            st.markdown("---")
            st.markdown("### Feedback")
            feedback = st.radio(
                "How helpful was this response?",
                ["Very helpful", "Somewhat helpful", "Not helpful"]
            )
            feedback_text = st.text_area("Additional feedback (optional):", height=100)
            if st.button("Submit Feedback"):
                st.success("Thank you for your feedback!")
            
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")
        logger.error(f"Error processing query: {e}", exc_info=True)
else:
    st.info("👆 Enter your question about hydrogen certifications above to get started!")

# Add sidebar with information
with st.sidebar:
    st.markdown("## About")
    st.markdown("""
    This tool helps answer questions about hydrogen certification standards using 
    a Retrieval-Augmented Generation (RAG) system.
    
    The system:
    1. Classifies which certification your question is about
    2. Optimizes your query
    3. Retrieves relevant information
    4. Generates a precise answer
    """)
    
    st.markdown("## Available Certifications")
    try:
        from app import list_certifications
        certifications = run_async(list_certifications())
        for cert in certifications:
            st.markdown(f"- {cert}")
    except Exception as e:
        st.warning("Could not load certification list")