import gradio as gr
import json
from urllib.parse import parse_qs, urlparse

# Function to process messages with context
def process_message(message, history, context_json):
    # Parse context from JSON string
    try:
        context = json.loads(context_json)
    except:
        context = {"error": "Invalid context data"}
    
    # Here you would connect to your chatbot backend
    # and provide both the message and the context
    
    # For demo purposes, we'll just echo the message with context info
    response = f"Received: {message}\n\nContext summary: Page titled '{context.get('title', 'Unknown')}'"
    
    # Format response for Gradio Chatbot using the newer messages format
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": response})
    return history

# Create Gradio interface
with gr.Blocks(css="#chatbot {height: 400px; overflow-y: auto;}") as demo:
    gr.HTML("<h2>Contextual Chatbot</h2>")
    
    # Hidden context state - will be updated via JavaScript
    context_json = gr.State('{}')
    
    # Optional: Show context for debugging
    with gr.Accordion("Page Context", open=False):
        context_display = gr.JSON(value={})
        update_button = gr.Button("Request Full Context")
    
    # Main chatbot interface
    chatbot = gr.Chatbot(value=[], type="messages")
    msg = gr.Textbox(placeholder="Ask about this page...")
    clear = gr.Button("Clear")
    
    # Handle sending messages
    msg.submit(process_message, [msg, chatbot, context_json], [chatbot])
    
    # Clear chat history
    clear.click(lambda: [], None, chatbot, queue=False)
    
    # JavaScript to handle context
    gr.HTML("""
    <script>
    // Debug function
    function logDebug(message, data) {
        console.log(`[DEBUG] ${message}`, data);
    }
    
    // Function to update the context state in Gradio
    function updateContextState(contextData) {
        logDebug("Updating context state with:", contextData);
        
        // Find the hidden state element for context
        const stateElements = Array.from(document.querySelectorAll('input[type="hidden"]'));
        const contextStateElem = stateElements.find(el => el.name && el.name.includes('context_json'));
        
        if (contextStateElem) {
            // Update the hidden input value
            contextStateElem.value = JSON.stringify(contextData);
            
            // Create and dispatch a change event
            const event = new Event('change', { bubbles: true });
            contextStateElem.dispatchEvent(event);
            
            logDebug("Context state updated", contextStateElem.value);
            return true;
        } else {
            console.error("Could not find context state element");
            return false;
        }
    }
    
    // Function to update the visible JSON display
    function updateJsonDisplay(contextData) {
        logDebug("Updating JSON display with:", contextData);
        
        // Find the JSON component's textarea
        const jsonTextarea = document.querySelector('.json-component textarea');
        if (jsonTextarea) {
            jsonTextarea.value = JSON.stringify(contextData, null, 2);
            
            // Find and click the submit button to update the value
            const submitButton = jsonTextarea.closest('.json-component').querySelector('button');
            if (submitButton) {
                submitButton.click();
                logDebug("JSON display updated");
                return true;
            }
        }
        
        console.error("Could not update JSON display");
        return false;
    }
    
    // Function to request full context
    function requestFullContext() {
        logDebug("Requesting full context from parent");
        window.parent.postMessage({type: 'getFullContext'}, '*');
    }
    
    // Try to get context from URL parameters
    function getContextFromUrl() {
        const urlParams = new URLSearchParams(window.location.search);
        const contextParam = urlParams.get('context');
        if (contextParam) {
            try {
                const parsedContext = JSON.parse(decodeURIComponent(contextParam));
                logDebug("Got context from URL:", parsedContext);
                return parsedContext;
            } catch (e) {
                console.error("Failed to parse context from URL:", e);
            }
        }
        return null;
    }
    
    // Listen for messages from parent (extension)
    window.addEventListener('message', function(event) {
        logDebug("Received message from parent:", event.data);
        
        if (event.data.type === 'fullContext') {
            // We got context data from the parent
            const contextData = event.data.content;
            
            // Update both the hidden state and visible display
            updateContextState(contextData);
            updateJsonDisplay(contextData);
        }
    });
    
    // Initialize when the page is fully loaded
    document.addEventListener('DOMContentLoaded', function() {
        logDebug("DOM fully loaded");
        
        // Wait for Gradio to fully initialize
        setTimeout(function() {
            logDebug("Initializing context handling");
            
            // Set up the update button
            const updateButton = document.querySelector('button.update-button');
            if (updateButton) {
                updateButton.addEventListener('click', requestFullContext);
                logDebug("Update button handler attached");
            }
            
            // First try to get context from URL
            const urlContext = getContextFromUrl();
            if (urlContext) {
                updateContextState(urlContext);
                updateJsonDisplay(urlContext);
            } else {
                // If no URL context, request from parent
                requestFullContext();
            }
        }, 2000); // Wait longer for Gradio to initialize
    });
    </script>
    """)
    
    # Add class to update button for JavaScript selection
    update_button.elem_classes = ["update-button"]

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True)