KATEX_CSS_JS = """
<style>
.gradio-container {max-width: 1200px !important}
#chatbot {height: 600px; overflow-y: auto;}
#sources_info {margin-top: 10px; color: #666;}

/* Improved styles for equations */
.katex { font-size: 1.1em !important; }
.math-inline { background: #f8f9fa; padding: 2px 5px; border-radius: 4px; }
.math-display { background: #f8f9f9; margin: 10px 0; padding: 10px; border-radius: 5px; overflow-x: auto; text-align: center; }

/* Table styles */
table {
    border-collapse: collapse;
    width: 100%;
    margin: 15px 0;
    font-size: 0.9em;
}
table, th, td {
    border: 1px solid #ddd;
}
th, td {
    padding: 8px 12px;
    text-align: left;
}
th {
    background-color: #f2f2f2;
}
tr:nth-child(even) {
    background-color: #f9f9f9;
}
.table-container {
    overflow-x: auto;
    margin-top: 10px;
}
</style>

<!-- Loading KaTeX -->
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.8/dist/katex.min.css">
<script src="https://cdn.jsdelivr.net/npm/katex@0.16.8/dist/katex.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/katex@0.16.8/dist/contrib/auto-render.min.js"></script>

<script>
// Function to process math equations with KaTeX
function renderMathInElement(element) {
    if (!window.renderMathInElement) return;
    
    try {
        window.renderMathInElement(element, {
            delimiters: [
                {left: '$$', right: '$$', display: true},
                {left: '$', right: '$', display: false},
                {left: '\\\\(', right: '\\\\)', display: false},
                {left: '\\\\[', right: '\\\\]', display: true}
            ],
            throwOnError: false,
            trust: true,
            strict: false,
            macros: {
                "\\\\R": "\\\\mathbb{R}",
                "\\\\N": "\\\\mathbb{N}"
            }
        });
    } catch (e) {
        console.error("KaTeX rendering error:", e);
    }
}

// Function to fix and prepare text for LaTeX rendering
function prepareTextForLatex(text) {
    if (!text) return text;
    
    // Don't modify code blocks
    if (text.indexOf('<pre>') !== -1) {
        const parts = text.split(/<pre>|<\/pre>/);
        for (let i = 0; i < parts.length; i++) {
            // Only process odd-indexed parts (non-code)
            if (i % 2 === 0) {
                parts[i] = prepareLatexInText(parts[i]);
            }
        }
        return parts.join('');
    }
    
    return prepareLatexInText(text);
}

// Helper to process LaTeX in regular text
function prepareLatexInText(text) {
    // Make sure dollar signs used for math have proper spacing
    // First, protect existing well-formed math expressions
    text = text.replace(/(\\$\\$[^\\$]+\\$\\$)/g, '<protect>$1</protect>');  // protect display math
    text = text.replace(/(\\$[^\\$\\n]+\\$)/g, '<protect>$1</protect>');    // protect inline math
    
    // Fix common LaTeX formatting issues outside protected regions
    text = text.replace(/([^<]protect[^>]*)(\\$)([^\\s])/g, '$1$2 $3');  // Add space after $ if needed
    text = text.replace(/([^\\s])(\\$)([^<]protect[^>]*)/g, '$1 $2$3');  // Add space before $ if needed
    
    // Handle subscripts: transform x_1 into x_{1} for better LaTeX compatibility
    text = text.replace(/([a-zA-Z])_([0-9a-zA-Z])/g, '$1_{$2}');
    
    // Restore protected content
    text = text.replace(/<protect>(.*?)<\/protect>/g, '$1');
    
    return text;
}

// Enhanced message processor for KaTeX rendering
function processMessage(message) {
    if (!message) return;
    
    try {
        // Get direct textual content when possible
        const elements = message.querySelectorAll('p, li, h1, h2, h3, h4, h5, span');
        elements.forEach(el => {
            const originalText = el.innerHTML;
            const preparedText = prepareTextForLatex(originalText);
            
            // Only update if changes were made
            if (preparedText !== originalText) {
                el.innerHTML = preparedText;
            }
            
            // Render equations in this element
            renderMathInElement(el);
        });
        
        // Also try to render on the entire message as fallback
        renderMathInElement(message);
    } catch (e) {
        console.error("Error processing message for LaTeX:", e);
    }
}

// Function to monitor for new messages
function setupMathObserver() {
    const chatElement = document.getElementById('chatbot');
    if (!chatElement) {
        setTimeout(setupMathObserver, 500);
        return;
    }
    
    // Process any existing messages
    chatElement.querySelectorAll('.message').forEach(processMessage);
    
    // Set up observer for new content
    const observer = new MutationObserver((mutations) => {
        for (const mutation of mutations) {
            if (mutation.addedNodes.length > 0 || mutation.type === 'characterData') {
                chatElement.querySelectorAll('.message').forEach(processMessage);
                break;
            }
        }
    });
    
    observer.observe(chatElement, { 
        childList: true, 
        subtree: true, 
        characterData: true 
    });
    
    console.log("LaTeX rendering observer set up successfully");
}

// Initialize once the document is fully loaded
function initializeRendering() {
    if (window.renderMathInElement) {
        setupMathObserver();
    } else {
        // If KaTeX isn't loaded yet, wait for it
        const katexScript = document.querySelector('script[src*="auto-render.min.js"]');
        if (katexScript) {
            katexScript.onload = setupMathObserver;
        } else {
            // Last resort: try again later
            setTimeout(initializeRendering, 500);
        }
    }
}

// Set up multiple trigger points to ensure it loads
document.addEventListener('DOMContentLoaded', function() {
    setTimeout(initializeRendering, 800);
});

window.addEventListener('load', function() {
    setTimeout(initializeRendering, 1200);
});
</script>
"""