import gradio as gr
from chromadb_semantic_search_for_dataset import search_cases, chat_with_cases

# Global variables to store current context
current_cases_context = ""
current_user_case = ""

def search_and_store_cases(query):
    """Search for cases and store context for chat"""
    global current_cases_context, current_user_case
    if not query.strip():
        return "कृपया केस खोज्नका लागि केही लेख्नुहोस्।", ""
    
    formatted_results, context, user_case = search_cases(query)
    current_cases_context = context
    current_user_case = user_case
    return formatted_results, ""

def chat_function(message, history):
    """Main chat function that uses current case context"""
    global current_cases_context, current_user_case
    
    if not message.strip():
        return history, ""
    
    # Generate response using current case context and user's case
    bot_response = chat_with_cases(message, current_cases_context, current_user_case)
    
    # Add to chat history
    history.append([message, bot_response])
    return history, ""

def clear_context():
    """Clear the stored case context"""
    global current_cases_context, current_user_case
    current_cases_context = ""
    current_user_case = ""
    return "केस कन्टेक्स्ट सफा गरियो। अब तपाईं सामान्य च्याट गर्न सक्नुहुन्छ।", [], ""

# Custom CSS
css = """
.gradio-container {
    max-width: 1200px !important;
    margin: auto !important;
}
.main-header {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 20px;
    border-radius: 10px;
    text-align: center;
    margin-bottom: 20px;
}
.section-header {
    background: #f8f9fa;
    padding: 15px;
    border-radius: 8px;
    border-left: 4px solid #007bff;
    margin: 10px 0;
}
.chat-container {
    border: 2px solid #e9ecef;
    border-radius: 10px;
    padding: 10px;
}
.search-container {
    border: 2px solid #e9ecef;
    border-radius: 10px;
    padding: 10px;
    background: #f8f9fa;
}
"""

# Create the main interface
with gr.Blocks(css=css, title="नेपाली कानूनी सहायक - Legal RAG System") as demo:
    
    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>⚖️ नेपाली कानूनी सहायक</h1>
        <h2>Nepali Legal Assistant with RAG</h2>
        <p><strong>आफ्नो मुद्दा खोज्नुहोस् र कानूनी सहायताका लागि च्याट गर्नुहोस्</strong></p>
        <p><em>Search for your case and chat for legal assistance</em></p>
    </div>
    """)
    
    with gr.Row():
        # Left Column - Case Search
        with gr.Column(scale=1):
            gr.HTML('<div class="section-header"><h3 style="color: #333;">🔍 केस खोज्नुहोस् (Case Search)</h3></div>')
            
            with gr.Group():
                search_input = gr.Textbox(
                    label="आफ्नो मुद्दाको विवरण लेख्नुहोस्",
                    placeholder="जस्तै: सम्पत्ति विवाद, हत्या मुद्दा, आर्थिक कसूर...",
                    lines=3
                )
                
                with gr.Row():
                    search_btn = gr.Button("🔍 खोज्नुहोस्", variant="primary", scale=2)
                    clear_btn = gr.Button("🗑️ सफा गर्नुहोस्", variant="secondary", scale=1)
                
                search_results = gr.Textbox(
                    label="समान मुद्दाहरू (Similar Cases)",
                    lines=12,
                    max_lines=15,
                    interactive=False
                )
                
                gr.Examples(
                    examples=[
                        ["जमीन जग्गाको विवाद"],
                        ["घरेलु हिंसा"],
                        ["ऋण असुली"],
                        ["हत्या मुद्दा"],
                        ["आर्थिक ठगी"],
                        ["चोरी मुद्दा"],
                        ["बलात्कार मुद्दा"],
                        ["भ्रष्टाचार मुद्दा"]
                    ],
                    inputs=search_input,
                    label="उदाहरणहरू (Examples)"
                )
        
        # Right Column - Chat Interface  
        with gr.Column(scale=1):
            gr.HTML('<div class="section-header"><h3 style="color: #333;">💬 कानूनी सहायता च्याट (Legal Chat)</h3></div>')
            
            with gr.Group():
                chatbot_ui = gr.Chatbot(
                    label="च्याट",
                    height=400,
                    show_copy_button=True,
                    bubble_full_width=False
                )
                
                with gr.Row():
                    msg_input = gr.Textbox(
                        placeholder="आफ्नो प्रश्न सोध्नुहोस्...",
                        show_label=False,
                        scale=4,
                        lines=1
                    )
                    send_btn = gr.Button("📤 पठाउनुहोस्", scale=1, variant="primary")
                
                with gr.Row():
                    clear_chat_btn = gr.Button("🗑️ च्याट सफा गर्नुहोस्", variant="secondary", scale=1)
                    clear_context_btn = gr.Button("🔄 कन्टेक्स्ट रिसेट गर्नुहोस्", variant="secondary", scale=1)
                
                context_status = gr.Textbox(
                    label="स्थिति (Status)",
                    value="सामान्य च्याट मोड - कुनै केस कन्टेक्स्ट छैन",
                    interactive=False,
                    lines=1
                )
                
                gr.Examples(
                    examples=[
                        ["मेरो केसमा के फैसला हुन सक्छ?"],
                        ["यी केसहरूसँग मेरो मुद्दाको समानता के छ?"],
                        ["मेरो केसको भिन्नता के छ?"],
                        ["यस्ता मुद्दामा कति समय लाग्छ?"],
                        ["के प्रमाणहरू चाहिन्छ?"],
                        ["मेरो केसमा जित्ने सम्भावना कति छ?"],
                        ["यी फैसलाहरूबाट के सिक्न सकिन्छ?"],
                        ["अदालती प्रक्रिया कस्तो हुन्छ?"]
                    ],
                    inputs=msg_input,
                    label="च्याट उदाहरणहरू (Chat Examples)"
                )
    
    # Information Section
    gr.HTML("""
    <div style="background: #e8f4fd; padding: 20px; border-radius: 10px; margin: 20px 0;">
        <h3 style="color: #0066cc;">📋 कसरी प्रयोग गर्ने (How to Use)</h3>
        <ol style="color: #333;">
            <li style="color: #333;"><strong style="color: #333;">केस खोज्नुहोस्:</strong> बाँया भागमा आफ्नो मुद्दाको विवरण लेखेर खोज्नुहोस्</li>
            <li style="color: #333;"><strong style="color: #333;">परिणाम हेर्नुहोस्:</strong> ५ वटा समान मुद्दाहरू देखाउनेछ</li>
            <li style="color: #333;"><strong style="color: #333;">च्याट गर्नुहोस्:</strong> दाँया भागमा ती मुद्दाहरूको बारेमा प्रश्न सोध्नुहोस्</li>
            <li style="color: #333;"><strong style="color: #333;">सामान्य च्याट:</strong> केस नखोजीकन पनि कानूनी प्रश्न सोध्न सक्नुहुन्छ</li>
        </ol>
        <p style="color: #d32f2f; font-weight: bold;">
            ⚠️ चेतावनी: यो केवल जानकारीमूलक सहायता हो। वास्तविक कानूनी सल्लाहका लागि योग्य वकिलसँग सल्लाह लिनुहोस्।
        </p>
    </div>
    """)
    
    # Event handlers
    search_btn.click(
        fn=search_and_store_cases,
        inputs=[search_input],
        outputs=[search_results, context_status]
    ).then(
        fn=lambda: "तपाईंको केस र समान केसहरूको कन्टेक्स्ट लोड भयो। अब तुलनाका लागि च्याट गर्न सक्नुहुन्छ।",
        outputs=[context_status]
    )
    
    clear_btn.click(
        fn=lambda: ("", ""),
        outputs=[search_input, search_results]
    )
    
    msg_input.submit(
        chat_function,
        inputs=[msg_input, chatbot_ui],
        outputs=[chatbot_ui, msg_input]
    )
    
    send_btn.click(
        chat_function,
        inputs=[msg_input, chatbot_ui],
        outputs=[chatbot_ui, msg_input]
    )
    
    clear_chat_btn.click(
        lambda: [],
        outputs=[chatbot_ui]
    )
    
    clear_context_btn.click(
        clear_context,
        outputs=[context_status, chatbot_ui, msg_input]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )