import streamlit as st
from qa_backend import DogFoodQASystem
import time
from typing import Dict, Any, List
import logging

import nltk
nltk.download('punkt_tab')

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Configure page settings
st.set_page_config(
    page_title="Dog Food Advisor",
    page_icon="🐕",
    layout="wide"
)

# Custom CSS for better styling
st.markdown("""
    <style>
    .stAlert {
        padding: 1rem;
        margin: 1rem 0;
        border-radius: 0.5rem;
    }
    .search-result {
        padding: 1rem;
        margin: 0.5rem 0;
        border: 1px solid #ddd;
        border-radius: 0.5rem;
    }
    .debug-info {
        font-size: small;
        color: gray;
        padding: 0.5rem;
        background-color: #f0f0f0;
        border-radius: 0.3rem;
    }
    </style>
    """, unsafe_allow_html=True)

@st.cache_resource
def load_qa_system() -> DogFoodQASystem:
    """Initialize and cache the QA system."""
    qa_system = DogFoodQASystem()
    # Run diagnostics
    vector_store_status = qa_system.diagnose_vector_store()
    return qa_system, vector_store_status

def display_search_result(result: Dict[str, Any], index: int) -> None:
    """Display a single search result with enhanced source and score information."""
    with st.container():
        # Source indicator and styling
        sources = result['sources']
        if len(sources) > 1:
            source_color = "#9C27B0"  # Purple for both sources
            source_badge = "🔄 Found in Both Sources"
            scores_text = f"BM25: {result['original_scores']['BM25']:.3f}, Vector: {result['original_scores']['Vector']:.3f}"
        elif 'Vector' in sources:
            source_color = "#2E7D32"  # Green for Vector
            source_badge = "🟢 Vector Search"
            scores_text = f"Score: {result['original_scores']['Vector']:.3f}"
        else:
            source_color = "#1565C0"  # Blue for BM25
            source_badge = "🔵 BM25 Search"
            scores_text = f"Score: {result['original_scores']['BM25']:.3f}"
        
        # Display header with source and score information
        st.markdown(f"""
        <div class="search-result">
            <h4 style="color: {source_color}">
                Result {index + 1} | {source_badge} | {scores_text}
            </h4>
        </div>
        """, unsafe_allow_html=True)
        
        # Display product details
        col1, col2 = st.columns(2)
        with col1:
            st.write("**Product Details:**")
            st.write(f"• Brand: {result['metadata']['brand']}")
            st.write(f"• Product: {result['metadata']['product_name']}")
            st.write(f"• Price: ${result['metadata']['price']:.2f}")
        
        with col2:
            st.write("**Additional Information:**")
            st.write(f"• Weight: {result['metadata']['weight']}kg")
            st.write(f"• Dog Type: {result['metadata']['dog_type']}")
            if 'reviews' in result['metadata']:
                st.write(f"• Reviews: {result['metadata']['reviews']}")
        
        st.markdown("**Description:**")
        st.write(result['text'])
        st.markdown("---")

def display_search_stats(results: List[Dict[str, Any]]) -> None:
    """Display detailed statistics about search results."""
    total_results = len(results)
    duplicates = sum(1 for r in results if len(r['sources']) > 1)
    vector_only = sum(1 for r in results if r['sources'] == ['Vector'])
    bm25_only = sum(1 for r in results if r['sources'] == ['BM25'])
    
    st.markdown("#### Search Results Statistics")
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.metric("Total Unique Results", total_results)
    with col2:
        st.metric("Found in Both Sources", duplicates, "🔄")
    with col3:
        st.metric("Vector Only", vector_only, "🟢")
    with col4:
        st.metric("BM25 Only", bm25_only, "🔵")

def main():
    # Header
    st.title("🐕 Dog Food Advisor")
    st.markdown("""
    Ask questions about dog food products in English or Spanish. 
    The system will provide relevant recommendations based on your query.
    """)
    
    # Initialize QA system with diagnostics
    qa_system, vector_store_status = load_qa_system()
    
    # Display system status
    with st.sidebar:
        st.markdown("### System Status")
        if vector_store_status:
            st.success("Vector Store: Connected")
        else:
            st.error("Vector Store: Not Connected")
            st.warning("Only BM25 search will be available")
    
    # Query input
    query = st.text_input(
        "Enter your question:",
        placeholder="e.g., 'What's the best food for puppies?' or '¿Cuál es la mejor comida para perros adultos?'"
    )
    
    # Add a search button
    search_button = st.button("Search")
    
    if query and search_button:
        with st.spinner("Processing your query..."):
            try:
                # Process query
                start_time = time.time()
                result = qa_system.process_query(query)
                processing_time = time.time() - start_time
                
                # Display answer
                st.markdown("### Answer")
                st.write(result["answer"])
                
                # Display search stats
                display_search_stats(result["search_results"])
                
                # Display processing information
                st.markdown(f"""
                <div class='debug-info'>
                    Language detected: {result['language']} | 
                    Processing time: {processing_time:.2f}s
                </div>
                """, unsafe_allow_html=True)
                
                # Display search results in an expander
                with st.expander("View Relevant Products", expanded=False):
                    st.markdown("### Search Results")
                    for i, search_result in enumerate(result["search_results"]):
                        display_search_result(search_result, i)
                        
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")
                logging.error(f"Error processing query: {str(e)}", exc_info=True)
    
    # Add footer with instructions
    st.markdown("---")
    with st.expander("Usage Tips"):
        st.markdown("""
        - Ask questions in English or Spanish
        - Be specific about your dog's needs (age, size, special requirements)
        - Include price preferences (e.g., 'affordable', 'premium')
        - Results are ranked by relevance and include price, brand, and product details
        - Results are color-coded: 
          - 🔵 Blue: BM25 Search Results
          - 🟢 Green: Vector Search Results
        """)

if __name__ == "__main__":
    main()