# Streamlit App for Multimodal RAG System - Deployment Optimized
import os, pathlib
cache_root = pathlib.Path("/tmp")
(os.makedirs := None)  # placeholder to satisfy linters if needed
(pathlib.Path("/tmp/hf")).mkdir(parents=True, exist_ok=True)
(pathlib.Path("/tmp/torch")).mkdir(parents=True, exist_ok=True)
os.environ.setdefault("HF_HOME", "/tmp/hf")
os.environ.setdefault("TRANSFORMERS_CACHE", "/tmp/hf")
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp/hf")
os.environ.setdefault("HF_DATASETS_CACHE", "/tmp/hf")
os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", "/tmp/hf")
os.environ.setdefault("XDG_CACHE_HOME", "/tmp")
os.environ.setdefault("TORCH_HOME", "/tmp/torch")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


import streamlit as st
import os
import tempfile
from pathlib import Path
from config import (
    STREAMLIT_CONFIG, SUPPORTED_DOCUMENT_TYPES, 
    SUPPORTED_IMAGE_TYPES, SUPPORTED_AUDIO_TYPES, 
    MODELS_CONFIG, MAX_FILE_SIZE_MB, DEMO_ENABLED
)
from multimodal_rag import MultimodalRAG
from device_setup import get_gpu_memory_info, clear_gpu_memory

def create_demo_content():
    '''Create demo content for the system'''
    st.markdown("### 📋 Try These Examples:")

    col1, col2, col3 = st.columns(3)

    with col1:
        st.markdown("**📄 Documents**")
        st.markdown("- Upload PDF files")
        st.markdown("- Upload DOCX files")
        st.markdown("- Upload TXT files")

    with col2:
        st.markdown("**🖼️ Images**")
        st.markdown("- Upload PNG/JPG images")
        st.markdown("- Screenshots")
        st.markdown("- Diagrams & charts")

    with col3:
        st.markdown("**🎵 Audio**")
        st.markdown("- Upload MP3/WAV files")
        st.markdown("- Voice recordings")
        st.markdown("- Lecture audio")

def display_system_info(rag_system):
    '''Display system information in sidebar'''
    with st.sidebar:
        st.markdown("### 🔧 System Status")

        # Device info
        if rag_system.has_gpu:
            st.success("🚀 GPU Acceleration Enabled")
            try:
                gpu_info = get_gpu_memory_info()
                if gpu_info:
                    st.text(f"GPU Memory: {gpu_info}")
            except:
                pass
        else:
            st.info("💻 Running on CPU")

        # Model info
        st.markdown("### 🤖 Loaded Models")
        model_info = {
            "Text Encoder": MODELS_CONFIG["text_encoder"],
            "Vision Model": "CLIP ViT-Base",
            "Speech Model": "Whisper Base",
            "LLM": "GPT4All" if rag_system.llm else "Not Available"
        }

        for name, model in model_info.items():
            status = "✅" if model != "Not Available" else "❌"
            st.text(f"{status} {name}: {model}")

def main():
    '''Main Streamlit interface'''

    # Configure page
    st.set_page_config(
        page_title=STREAMLIT_CONFIG["page_title"],
        layout=STREAMLIT_CONFIG["layout"],
        page_icon="🤖",
        initial_sidebar_state=STREAMLIT_CONFIG["initial_sidebar_state"]
    )

    # Header
    st.title("🤖 Multimodal RAG System")
    st.markdown("### Upload documents, images, and audio files for AI-powered semantic search!")

    # Initialize RAG system
    if 'rag_system' not in st.session_state:
        with st.spinner("🔄 Loading AI models... This may take a few minutes on first run."):
            try:
                st.session_state.rag_system = MultimodalRAG()
                st.balloons()
                st.success("🎉 System loaded successfully!")
            except Exception as e:
                st.error(f"❌ Failed to load system: {e}")
                st.info("💡 Please refresh the page and try again")
                st.stop()

    rag = st.session_state.rag_system

    # Display system info in sidebar
    display_system_info(rag)

    # Main interface
    tab1, tab2, tab3 = st.tabs(["📤 Upload & Search", "📊 Knowledge Base", "ℹ️ About"])

    with tab1:
        # File upload section
        st.markdown("### 📁 Upload Files")

        uploaded_files = st.file_uploader(
            "Choose files to add to knowledge base",
            type=SUPPORTED_DOCUMENT_TYPES + SUPPORTED_IMAGE_TYPES + SUPPORTED_AUDIO_TYPES,
            accept_multiple_files=True,
            help=f"Supported formats: {', '.join(SUPPORTED_DOCUMENT_TYPES + SUPPORTED_IMAGE_TYPES + SUPPORTED_AUDIO_TYPES)}"
        )

        # Process uploaded files
        if uploaded_files:
            if st.button("🔄 Process Files", type="primary"):
                processed_count = 0
                failed_count = 0

                progress_bar = st.progress(0)
                status_text = st.empty()

                for i, uploaded_file in enumerate(uploaded_files):
                    try:
                        # Check file size
                        file_size_mb = uploaded_file.size / (1024 * 1024)
                        if file_size_mb > MAX_FILE_SIZE_MB:
                            st.warning(f"⚠️ {uploaded_file.name} is too large ({file_size_mb:.1f}MB > {MAX_FILE_SIZE_MB}MB)")
                            failed_count += 1
                            continue

                        status_text.text(f"Processing {uploaded_file.name}...")

                        # Save temporary file
                        with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp_file:
                            tmp_file.write(uploaded_file.getbuffer())
                            tmp_file_path = tmp_file.name

                        # Process file
                        success = rag.add_document(tmp_file_path, uploaded_file.name)

                        if success:
                            processed_count += 1
                            st.success(f"✅ Processed: {uploaded_file.name}")
                        else:
                            failed_count += 1
                            st.error(f"❌ Failed: {uploaded_file.name}")

                        # Clean up
                        os.unlink(tmp_file_path)

                    except Exception as e:
                        st.error(f"❌ Error processing {uploaded_file.name}: {e}")
                        failed_count += 1

                    progress_bar.progress((i + 1) / len(uploaded_files))

                status_text.text("Processing complete!")
                st.info(f"📊 Processed: {processed_count} files, Failed: {failed_count} files")

        # Search section
        st.markdown("### 🔍 Search Knowledge Base")

        query = st.text_input(
            "Enter your question or search query:",
            placeholder="e.g., What is the main topic of the uploaded document?",
            help="Ask questions about your uploaded content"
        )

        col1, col2 = st.columns([3, 1])

        with col1:
            top_k = st.slider("Number of results", 1, 10, 5)

        with col2:
            search_button = st.button("🔍 Search", type="primary")

        if search_button and query:
            with st.spinner("🔄 Searching..."):
                try:
                    results = rag.search(query, top_k=top_k)

                    if results:
                        st.markdown("### 📋 Search Results")

                        for i, result in enumerate(results, 1):
                            with st.expander(f"Result {i}: {result['filename']} (Score: {result['score']:.3f})"):
                                st.markdown(f"**Content:** {result['content']}")
                                st.markdown(f"**Type:** {result['content_type']}")

                        # Generate answer if LLM is available
                        if rag.llm:
                            st.markdown("### 🤖 AI Generated Answer")
                            with st.spinner("🧠 Generating answer..."):
                                try:
                                    answer = rag.generate_answer(query, results)
                                    st.markdown(answer)
                                except Exception as e:
                                    st.error(f"❌ Error generating answer: {e}")
                        else:
                            st.info("💡 Upload a question to get AI-generated answers (LLM not available)")

                    else:
                        st.warning("🤔 No results found. Try uploading some files first or use different keywords.")

                except Exception as e:
                    st.error(f"❌ Search failed: {e}")

    with tab2:
        # Knowledge base overview
        st.markdown("### 📚 Knowledge Base Overview")

        try:
            doc_count = len(rag.documents)
            st.metric("Total Documents", doc_count)

            if doc_count > 0:
                st.markdown("### 📄 Stored Documents")
                for i, doc in enumerate(rag.documents[:10]):  # Show first 10
                    st.text(f"• {doc['filename']} ({doc['content_type']})")

                if doc_count > 10:
                    st.text(f"... and {doc_count - 10} more documents")

                # Clear database option
                if st.button("🗑️ Clear Knowledge Base", help="Remove all documents"):
                    if st.confirm("Are you sure you want to clear all documents?"):
                        rag.clear_database()
                        st.success("✅ Knowledge base cleared!")
                        st.rerun()
            else:
                st.info("📭 No documents in knowledge base. Upload some files to get started!")

                if DEMO_ENABLED:
                    create_demo_content()

        except Exception as e:
            st.error(f"❌ Error accessing knowledge base: {e}")

    with tab3:
        # About section
        st.markdown("### ℹ️ About This System")

        st.markdown('''
        This **Multimodal RAG (Retrieval-Augmented Generation)** system allows you to:

        - 📄 **Process Documents**: Upload PDFs, DOCX, and text files
        - 🖼️ **Analyze Images**: Extract text and understand visual content  
        - 🎵 **Transcribe Audio**: Convert speech to searchable text
        - 🔍 **Semantic Search**: Find relevant content using AI embeddings
        - 🤖 **AI Answers**: Get intelligent responses from a local LLM

        ### 🛠️ Technology Stack
        - **Frontend**: Streamlit
        - **Embeddings**: Sentence Transformers + CLIP
        - **Vector Search**: FAISS
        - **Speech Recognition**: OpenAI Whisper
        - **OCR**: EasyOCR  
        - **Local LLM**: GPT4All

        ### 🚀 Features
        - ⚡ **GPU Acceleration** (when available)
        - 🔒 **Privacy First** (all processing runs locally/on your server)
        - 📱 **Responsive Design** 
        - 🎯 **Easy to Use**
        ''')

        # System stats
        st.markdown("### 📊 System Information")
        col1, col2 = st.columns(2)

        with col1:
            st.markdown("**Supported File Types:**")
            st.text(f"Documents: {', '.join(SUPPORTED_DOCUMENT_TYPES)}")
            st.text(f"Images: {', '.join(SUPPORTED_IMAGE_TYPES)}")
            st.text(f"Audio: {', '.join(SUPPORTED_AUDIO_TYPES)}")

        with col2:
            st.markdown("**Current Configuration:**")
            st.text(f"Max file size: {MAX_FILE_SIZE_MB}MB")
            st.text(f"Device: {'GPU' if rag.has_gpu else 'CPU'}")
            st.text(f"Models loaded: {len([m for m in [rag.text_encoder, rag.clip_model, rag.whisper_model, rag.ocr_reader, rag.llm] if m is not None])}/5")

if __name__ == "__main__":
    main()