"""
UI components and helper functions for the Defensor application.
"""

import gradio as gr
import time
import traceback
import uuid
import re

def create_fallback_ui(custom_css):
    """Create a fallback UI for when data loading fails."""
    with gr.Blocks(
        theme=gr.themes.Base(),
        css=custom_css,
        title="Defensor Legal Assistant - Data Error"
    ) as demo:
        # Create a header with app-header class for better styling
        with gr.Column(elem_classes="app-header"):
            gr.HTML("""
            <div class="logo-container">
                <div class="logo">
                    <span class="logo-icon">⚖️</span>
                    <span class="logo-text">Defensor</span>
                </div>
            </div>
            """)
            gr.Markdown("# Agentic Defensor - Data Loading Error")
        
        with gr.Column(elem_classes="main-container"):
            gr.Markdown("""
            ### There was an error loading the required data files.
            
            This application requires specific dataset files to function properly. 
            The system tried to download them from Hugging Face Hub but encountered errors.
            
            Please check the application logs for more details on what went wrong.
            
            Common issues:
            1. Dataset repository might not be accessible or might not contain the expected files
            2. Network connectivity issues when downloading datasets
            3. Missing or invalid Hugging Face token for accessing private datasets
            """)
            
            with gr.Row():
                input_text = gr.Textbox(label="You can enter text here, but functionality is limited")
                output_text = gr.Textbox(label="System response")
                
            submit_btn = gr.Button("Submit (Limited functionality)")
            
            def dummy_response(text):
                return "The system is in fallback mode due to data loading errors. Please check application logs."
                
            submit_btn.click(dummy_response, inputs=[input_text], outputs=[output_text])
        
    return demo

def create_interface(custom_css, process_query_function):
    """Create the main Gradio interface."""
    
    with gr.Blocks(
        title="Defensor Legal Assistant",
        theme=gr.themes.Base(),
        css=custom_css
    ) as demo:
        # Create a header with app-header class for better styling
        with gr.Column(elem_classes="app-header"):
            gr.HTML("""
            <div class="logo-container">
                <div class="logo">
                    <span class="logo-icon">⚖️</span>
                    <span class="logo-text">Defensor</span>
                </div>
            </div>
            """)
            gr.Markdown("# Defensor Legal Assistant")
            gr.Markdown("Ask questions about legal defense strategies and get comprehensive answers based on legitimate sources.")
        
        # Main application container
        with gr.Column(elem_classes="main-container"):
            # Add status indicator with styles - in its own container
            with gr.Row(elem_classes="status-container"):
                status_indicator = gr.HTML(
                    value="<div class='status-ready'>Ready to assist</div>",
                    label=None
                )
            
            # Enhanced container for the input area
            with gr.Column(elem_classes="input-container"):
                query_input = gr.Textbox(
                    label="Your Legal Question",
                    placeholder="Enter your legal question here (e.g., 'What are the key elements of self-defense?')",
                    lines=3
                )
                gr.Markdown("<div class='hint-text'>Ask detailed questions for better results. Be specific about the legal concepts you're interested in.</div>")
            
            # Create a dedicated container for the progress area
            with gr.Column(elem_classes="progress-container"):
                # Progress indicator with animation and better styling
                progress = gr.Progress()
                
                # Add a visual progress label
                progress_label = gr.HTML(
                    "<div class='progress-label'>Waiting for your question...</div>", 
                    visible=True
                )
            
            # Better button container with centered buttons
            with gr.Row(elem_classes="button-container"):
                submit_btn = gr.Button(
                    "Submit Question", 
                    variant="primary", 
                    interactive=True
                )
                clear_btn = gr.Button(
                    "Clear", 
                    variant="secondary"
                )
            
            # Settings in collapsible section
            with gr.Accordion("Settings & Options", open=False):
                with gr.Row():
                    top_k = gr.Slider(
                        minimum=5, 
                        maximum=50, 
                        value=20, 
                        step=5, 
                        label="Number of Documents to Consider"
                    )
                    debug = gr.Checkbox(label="Show Agent Reasoning Process", value=True)
            
            # Display area for the response with enhanced formatting
            with gr.Column(elem_classes="output-container"):
                output = gr.HTML(label="Response", value="")
        
        # About section at the bottom
        with gr.Accordion("About Defensor Legal Assistant", open=False):
            gr.Markdown("""
            ## About Defensor Legal Assistant
            
            Defensor is a specialized legal assistant powered by advanced AI technology and a comprehensive legal knowledge base. It provides insights and information about legal defense strategies based on authoritative sources.
            
            ### How It Works
            
            1. **Query Analysis**: The system analyzes your legal question to understand its context and intent
            2. **Knowledge Retrieval**: It searches through a database of legal documents to find relevant information
            3. **Context Organization**: The system organizes the retrieved information into a coherent context
            4. **Response Generation**: A comprehensive answer is generated based on the collected information
            
            ### Tips for Best Results
            
            - Be specific in your questions
            - Include relevant legal terms or concepts
            - For complex topics, consider breaking down into multiple targeted questions
            - Use the "Show Agent Reasoning" option to see how the system arrived at its answer
            
            ### Limitations
            
            This assistant provides legal information, not legal advice. Always consult with a qualified attorney for advice on specific legal matters.
            """)

        # Define status update helper functions
        def set_ready_status():
            return "<div class='status-ready'>Ready to assist</div>", "<div class='progress-label'>Waiting for your question...</div>"
        
        def set_processing_status():
            return "<div class='status-processing'>Processing your question...</div>", gr.update(interactive=False), "<div class='progress-label'>Processing legal query...</div>"
        
        def set_completed_status():
            return "<div class='status-complete'>Response complete</div>", gr.update(interactive=True), "<div class='progress-label'>Response complete</div>"
        
        def clear_inputs():
            return "", "<div class='status-ready'>Ready to assist</div>", gr.update(interactive=True), "", "<div class='progress-label'>Waiting for your question...</div>"
        
        # Set up event handlers with improved button state management
        submit_btn.click(
            fn=set_processing_status,
            outputs=[status_indicator, submit_btn, progress_label],
            queue=False
        ).then(
            fn=process_query_function,
            inputs=[query_input, top_k, debug],
            outputs=[output, progress_label],
            api_name="query"
        ).then(
            fn=set_completed_status,
            outputs=[status_indicator, submit_btn, progress_label],
            queue=False
        )
        
        clear_btn.click(
            fn=clear_inputs,
            outputs=[query_input, status_indicator, submit_btn, output, progress_label],
            queue=False
        )
        
        return demo
        
def format_progress_stages(stage, is_error=False):
    """Helper to format the progress stages HTML."""
    if stage == 1:
        return """<div class='progress-stages-container'>
            <div class='progress-stage active'>► Analyzing query</div>
            <div class='progress-stage'>Searching documents</div>
            <div class='progress-stage'>Organizing information</div>
            <div class='progress-stage'>Formulating response</div>
        </div>"""
    elif stage == 2:
        return """<div class='progress-stages-container'>
            <div class='progress-stage complete'>✓ Analyzing query</div>
            <div class='progress-stage active'>► Searching documents</div>
            <div class='progress-stage'>Organizing information</div>
            <div class='progress-stage'>Formulating response</div>
        </div>"""
    elif stage == 3:
        return """<div class='progress-stages-container'>
            <div class='progress-stage complete'>✓ Analyzing query</div>
            <div class='progress-stage complete'>✓ Searching documents</div>
            <div class='progress-stage active'>► Organizing information</div>
            <div class='progress-stage'>Formulating response</div>
        </div>"""
    elif stage == 4:
        return """<div class='progress-stages-container'>
            <div class='progress-stage complete'>✓ Analyzing query</div>
            <div class='progress-stage complete'>✓ Searching documents</div>
            <div class='progress-stage complete'>✓ Organizing information</div>
            <div class='progress-stage active'>► Formulating response</div>
        </div>"""
    elif stage == 5:  # Complete
        return """<div class='progress-stages-container'>
            <div class='progress-stage complete'>✓ Analyzing query</div>
            <div class='progress-stage complete'>✓ Searching documents</div>
            <div class='progress-stage complete'>✓ Organizing information</div>
            <div class='progress-stage complete'>✓ Response ready</div>
        </div>"""
    elif is_error:  # Error state
        return """<div class='progress-stages-container'>
            <div class='progress-stage complete'>✓ Analyzing query</div>
            <div class='progress-stage complete'>✓ Searching documents</div>
            <div class='progress-stage error'>✗ Error</div>
            <div class='progress-stage'>Formulating response</div>
        </div>"""
    
    # Default fallback
    return ""

def format_answer_html(answer, sources=None, reasoning_steps=None, debug=False):
    """Format the answer into well-structured HTML with improved Markdown rendering."""
    
    # Process Markdown more comprehensively
    answer_html = ""
    
    # Process paragraphs and add proper formatting
    for paragraph in answer.split("\n\n"):
        if not paragraph.strip():
            continue
            
        # Process headers
        if paragraph.startswith("# "):
            heading = paragraph[2:].strip()
            answer_html += f"<h2>{heading}</h2>"
        elif paragraph.startswith("## "):
            heading = paragraph[3:].strip()
            answer_html += f"<h3>{heading}</h3>"
        elif paragraph.startswith("### "):
            heading = paragraph[4:].strip()
            answer_html += f"<h4>{heading}</h4>"
        # Process lists
        elif re.match(r'^\d+\.\s', paragraph):
            # Ordered list
            items = paragraph.split("\n")
            answer_html += "<ol>"
            for item in items:
                if re.match(r'^\d+\.\s', item):
                    item_content = re.sub(r'^\d+\.\s', '', item)
                    answer_html += f"<li>{item_content}</li>"
            answer_html += "</ol>"
        elif re.match(r'^[*-]\s', paragraph):
            # Unordered list
            items = paragraph.split("\n")
            answer_html += "<ul>"
            for item in items:
                if re.match(r'^[*-]\s', item):
                    item_content = re.sub(r'^[*-]\s', '', item)
                    answer_html += f"<li>{item_content}</li>"
            answer_html += "</ul>"
        # Process code blocks
        elif paragraph.startswith("```") and paragraph.endswith("```"):
            code = paragraph[3:-3].strip()
            language = ""
            if "\n" in code:
                first_line = code.split("\n")[0].strip()
                if first_line and not first_line.startswith("```"):
                    language = first_line
                    code = "\n".join(code.split("\n")[1:])
            
            answer_html += f'<pre class="code-block"><code class="language-{language}">{code}</code></pre>'
        # Process blockquotes
        elif paragraph.startswith(">"):
            content = paragraph.replace("\n>", "\n").replace(">", "")
            answer_html += f'<blockquote>{content}</blockquote>'
        # Standard paragraph
        else:
            # Process inline code
            paragraph = re.sub(r'`([^`]+)`', r'<code>\1</code>', paragraph)
            # Process bold text
            paragraph = re.sub(r'\*\*([^*]+)\*\*', r'<strong>\1</strong>', paragraph)
            # Process italic text
            paragraph = re.sub(r'\*([^*]+)\*', r'<em>\1</em>', paragraph)
            
            answer_html += f"<p>{paragraph}</p>"
    
    # Add sources if available
    if sources:
        sources_html = "<div class='sources-section'>"
        sources_html += "<h3>Sources</h3><ul>"
        for source in sources:
            if source and source != "Source information unavailable":
                sources_html += f"<li class='source-item'>{source}</li>"
        sources_html += "</ul></div>"
        answer_html += sources_html
    
    # Add debugging information if available with nice formatting
    if debug and reasoning_steps:
        reasoning_html = "<h3>Agent Reasoning Process</h3>"
        for step in reasoning_steps:
            step_stage = step['stage']
            step_reasoning = step['reasoning'].replace('\n', '<br>')
            reasoning_html += f"""
            <div class='reasoning-step'>
                <div class='reasoning-stage'>{step_stage}</div>
                <div>{step_reasoning}</div>
            </div>
            """
        answer_html += reasoning_html
    
    # Return the final formatted HTML without any copy functionality
    return f"""
    <div class="response-wrapper">
        <div class="answer-container">
            {answer_html}
        </div>
    </div>
    """