#!/usr/bin/env python3
"""
Usage: python scripts/launch_ui.py

Gradio web UI for the RAG system.
Provides a user-friendly interface for asking questions.
"""
import gradio as gr
from query.retriever import retrieve
from query.prompt_builder import build_prompt
from query.generator import generate


def answer_question_generator(question: str, show_sources: bool = False, tutor_mode: bool = False):
    """
    Answer a question using the RAG system with streaming status updates.
    
    Args:
        question: The user's question
        show_sources: Whether to display retrieved sources
        tutor_mode: If True, use Socratic tutor mode; if False, use direct answer mode
    
    Yields:
        Tuples of (answer_markdown, sources_markdown) for progressive UI updates
    """
    if not question.strip():
        yield "Please enter a question.", ""
        return
    
    # Step 1: Retrieving sources
    yield "🔍 **Retrieving relevant sources...**\n\n*(This may take 5-10 seconds)*", ""
    
    try:
        chunks = retrieve(question)
        count = len(chunks)
        
        mode_label = "🎓 **Tutor Mode**" if tutor_mode else "📚 **Regular Mode**"
        yield f"✓ **Retrieved {count} relevant sources**\n\n⏳ **Generating response ({mode_label})...**\n\n*(This may take 15-30 seconds with local Ollama. Please wait.)*", ""
        
        # Build prompt with context, passing tutor_mode
        prompt = build_prompt(question, chunks, tutor_mode=tutor_mode)
        
        # Generate response (this is the long operation)
        response = generate(prompt)
        
        # Format sources
        sources_text = ""
        if show_sources and chunks:
            sources_text = "## Retrieved Sources\n\n"
            for i, chunk in enumerate(chunks, 1):
                source = chunk["metadata"].get("source", "unknown")
                page = chunk["metadata"].get("page", chunk["metadata"].get("slide", "?"))
                category = chunk["metadata"].get("category", "unknown")
                distance = chunk.get("distance", "?")
                sources_text += f"### Source {i}\n"
                sources_text += f"**Category:** {category}\n"
                sources_text += f"**Location:** {source} (page/slide {page})\n"
                sources_text += f"**Relevance:** {distance:.4f}\n"
                sources_text += f"**Excerpt:** {chunk['text'][:200]}...\n\n"
        
        # Final: Show answer
        yield response, sources_text
        
    except Exception as e:
        yield f"❌ **Error:** {str(e)}", ""


def main():
    """Launch the Gradio interface."""
    with gr.Blocks(title="ENGR 445 RAG Assistant") as demo:
        gr.Markdown("# ENGR 445 Course Assistant")
        gr.Markdown("Ask questions about ENGR 445 course material — lectures, datasheets, application notes, and source code.")
        
        with gr.Row():
            question_input = gr.Textbox(
                label="Your Question",
                placeholder="E.g., How do I configure GPIO pins as push-pull?",
                lines=2
            )
        
        with gr.Row():
            submit_btn = gr.Button("Ask", scale=1)
            show_sources_cb = gr.Checkbox(label="Show Retrieved Sources", value=False)
            tutor_mode_cb = gr.Checkbox(label="🎓 Tutor Mode (hints instead of answers)", value=False)
        
        with gr.Row():
            answer_output = gr.Markdown(label="Answer", value="*Ask a question to get started...*")
        
        with gr.Row():
            sources_output = gr.Markdown(label="Sources")
        
        with gr.Row():
            gr.Markdown("""
**Mode Explanation:**
- **Regular Mode**: Provides direct answers with citations
- **Tutor Mode**: Acts as a Socratic tutor, providing hints and guiding you to resources instead of answers
            """)
        
        # Connect the submit button to the function with streaming updates
        submit_btn.click(
            fn=answer_question_generator,
            inputs=[question_input, show_sources_cb, tutor_mode_cb],
            outputs=[answer_output, sources_output]
        )
        
        # Also support Enter key
        question_input.submit(
            fn=answer_question_generator,
            inputs=[question_input, show_sources_cb, tutor_mode_cb],
            outputs=[answer_output, sources_output]
        )
        
        gr.Markdown("---")
        gr.Markdown("**Note:** All inference runs locally via Ollama. Ensure Ollama is running before using this interface.")
    
    demo.launch(server_name="127.0.0.1", server_port=7860, share=False)


if __name__ == "__main__":
    main()