import streamlit as st
from utils.rag_system import DocumentIngestion, RAGSystem, save_query_result, load_shared_query

def display_query_result(result, show_share_link=False):
    """Display query results in a formatted way."""
    st.markdown('<div class="query-result">', unsafe_allow_html=True)
    
    # Show which model was used
    if result.get("model_used"):
        st.info(f"🤖 **Model Used:** {result['model_used']}")
    
    st.subheader("🎯 Answer")
    st.write(result["answer"])
    
    # Share link
    if show_share_link and result.get("query_id"):
        st.markdown("---")
        current_url = st.get_option("browser.serverAddress") or "localhost:8501"
        share_url = f"http://{current_url}?share={result['query_id']}"
        st.markdown(f"""
        <div class="share-link">
            <strong>🔗 Share this result:</strong><br>
            <code>{share_url}</code>
        </div>
        """, unsafe_allow_html=True)
        
        if st.button("📋 Copy Share Link"):
            st.code(share_url)
    
    # Source documents
    if result.get("source_documents"):
        st.markdown("---")
        st.subheader("📚 Sources")
        for i, doc in enumerate(result["source_documents"], 1):
            with st.expander(f"Source {i}: {doc.metadata.get('source', 'Unknown')}"):
                col1, col2 = st.columns([1, 2])
                with col1:
                    st.write(f"**University:** {doc.metadata.get('university', 'Unknown')}")
                    st.write(f"**Country:** {doc.metadata.get('country', 'Unknown')}")
                    st.write(f"**Type:** {doc.metadata.get('document_type', 'Unknown')}")
                with col2:
                    st.write("**Relevant Content:**")
                    content_preview = doc.page_content[:300] + "..." if len(doc.page_content) > 300 else doc.page_content
                    st.write(content_preview)
    
    st.markdown('</div>', unsafe_allow_html=True)

def display_shared_query(query_id):
    """Display a shared query result."""
    st.header("🔗 Shared Query Result")
    
    result_data = load_shared_query(query_id)
    
    if result_data:
        st.info(f"**Original Question:** {result_data['question']}")
        st.write(f"**Language:** {result_data['language']}")
        st.write(f"**Date:** {result_data['timestamp'][:10]}")
        
        # Create a mock result object for display
        mock_result = {
            "answer": result_data["answer"],
            "source_documents": [
                type('MockDoc', (), {
                    'metadata': source,
                    'page_content': source.get('content_preview', '')
                })() for source in result_data.get('sources', [])
            ]
        }
        
        display_query_result(mock_result, show_share_link=False)
        
        if st.button("🔍 Ask Your Own Question"):
            st.experimental_set_query_params()
            st.experimental_rerun()
    else:
        st.error("❌ Shared query not found or has expired.")
        if st.button("🏠 Go to Home"):
            st.experimental_set_query_params()
            st.experimental_rerun()