Vimium-style Highlight-All Enhanced

// ==UserScript==
// @name         Highlight-All Enhanced
// @namespace    https://github.com/yourusername/highlight-all-enhanced
// @version      1.1.0
// @description  Enhanced highlight script with improved performance and error handling
// @author       Your Name
// @match        *://*/*
// @grant        none
// @run-at       document-idle
// ==/UserScript==

(function() {
    'use strict';

    // ====================================================================
    // ENHANCED STATE AND CONFIGURATION
    // ====================================================================

    let currentObserver = null;
    let currentPattern = null;
    let isHighlighting = false;
    let mutationThrottle = null;
    let pendingMutations = new Set();

    const HIGHLIGHT_CLASS = 'highlight-all-mark';
    const SKIP_TAGS = new Set(['SCRIPT', 'STYLE', 'TEMPLATE', 'NOSCRIPT', 'HEAD', 'TITLE', 'SVG']);
    const VALID_FLAGS = new Set(['d', 'g', 'i', 'm', 's', 'u', 'y']);
    const MAX_ZERO_LENGTH_ITERATIONS = 1000; // Safety limit for problematic patterns

    // Enhanced idle callback with better fallback
    const scheduleIdleWork = window.requestIdleCallback || function(callback, options = {}) {
        const timeout = options.timeout || 50;
        return setTimeout(() => {
            const start = performance.now();
            callback({
                timeRemaining: () => Math.max(0, 50 - (performance.now() - start)),
                didTimeout: false
            });
        }, timeout);
    };

    const cancelIdleWork = window.cancelIdleCallback || clearTimeout;

    // ====================================================================
    // ENHANCED PATTERN PARSING WITH STRICT VALIDATION
    // ====================================================================

    /**
     * Enhanced pattern parsing with strict flag validation and better error handling
     * @param {string} input - User input pattern
     * @returns {RegExp|null} Validated and compiled regex
     */
    function parsePattern(input) {
        if (!input || typeof input !== 'string') return null;

        input = input.trim();
        if (!input) return null;

        // Detect regex literal format: /pattern/flags
        const regexMatch = input.match(/^\/(.+?)\/([dgimsuyx]*)$/);

        if (regexMatch) {
            try {
                const [, pattern, flags] = regexMatch;

                // Validate flags - only allow valid JavaScript regex flags
                const validFlags = flags.split('').filter(f => VALID_FLAGS.has(f));
                const flagString = [...new Set(validFlags)].join(''); // Remove duplicates

                // Ensure global flag for "highlight all" behavior
                const globalFlags = flagString.includes('g') ? flagString : flagString + 'g';

                return new RegExp(pattern, globalFlags);
            } catch (error) {
                console.warn('Invalid regex pattern:', input, error.message);
                return null;
            }
        } else {
            // Plain string: escape metacharacters with enhanced escaping
            try {
                const escapedPattern = input.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
                return new RegExp(escapedPattern, 'gi');
            } catch (error) {
                console.warn('Failed to create regex from string:', input, error.message);
                return null;
            }
        }
    }

    // ====================================================================
    // ENHANCED DOM TRAVERSAL WITH IMPROVED FILTERING
    // ====================================================================

    /**
     * Enhanced TreeWalker filter with better skip logic and shadow DOM awareness
     * @param {Node} node - Text node to evaluate
     * @returns {number} NodeFilter result
     */
    function enhancedTextNodeFilter(node) {
        // Skip empty or whitespace-only text nodes
        if (!node.data || !node.data.trim()) {
            return NodeFilter.FILTER_REJECT;
        }

        // Check parent chain for problematic elements
        let parent = node.parentNode;
        while (parent && parent.nodeType === Node.ELEMENT_NODE) {
            // Skip if parent is problematic tag
            if (SKIP_TAGS.has(parent.tagName)) {
                return NodeFilter.FILTER_REJECT;
            }

            // Skip if already inside highlight mark
            if (parent.classList && parent.classList.contains(HIGHLIGHT_CLASS)) {
                return NodeFilter.FILTER_REJECT;
            }

            // Skip contenteditable="false" elements
            if (parent.contentEditable === 'false') {
                return NodeFilter.FILTER_REJECT;
            }

            parent = parent.parentNode;
        }

        return NodeFilter.FILTER_ACCEPT;
    }

    /**
     * Create enhanced TreeWalker with improved filtering
     * @param {Node} root - Root node to traverse
     * @returns {TreeWalker}
     */
    function createEnhancedTextWalker(root) {
        return document.createTreeWalker(
            root,
            NodeFilter.SHOW_TEXT,
            enhancedTextNodeFilter,
            false
        );
    }

    // ====================================================================
    // ENHANCED MATCHING WITH BETTER ZERO-LENGTH PROTECTION
    // ====================================================================

    /**
     * Enhanced text node wrapping with improved zero-length match handling
     * @param {Text} textNode - Text node to process
     * @param {RegExp} regex - Pattern to match
     */
    function enhancedWrapMatches(textNode, regex) {
        if (!textNode || !textNode.data || !textNode.parentNode) return;

        // Reset regex state
        regex.lastIndex = 0;

        let currentNode = textNode;
        let match;
        let iterationCount = 0;
        const maxIterations = Math.min(MAX_ZERO_LENGTH_ITERATIONS, currentNode.data.length * 2);

        while ((match = regex.exec(currentNode.data)) !== null && iterationCount < maxIterations) {
            iterationCount++;

            const matchStart = match.index;
            const matchLength = match[0].length;

            // Enhanced zero-length match protection
            if (matchLength === 0) {
                // If lastIndex didn't advance, manually increment to prevent infinite loop
                if (regex.lastIndex === match.index) {
                    regex.lastIndex = Math.min(match.index + 1, currentNode.data.length);
                }
                // Break if we've reached the end
                if (regex.lastIndex >= currentNode.data.length) break;
                continue;
            }

            // Validate match boundaries
            if (matchStart < 0 || matchStart >= currentNode.data.length) break;
            if (matchStart + matchLength > currentNode.data.length) break;

            try {
                // Split text node at exact match boundaries
                let matchNode;

                if (matchStart > 0) {
                    matchNode = currentNode.splitText(matchStart);
                    currentNode = matchNode;
                } else {
                    matchNode = currentNode;
                }

                let afterMatch;
                if (matchLength < matchNode.data.length) {
                    afterMatch = matchNode.splitText(matchLength);
                }

                // Create enhanced highlight mark
                const mark = document.createElement('mark');
                mark.className = HIGHLIGHT_CLASS;
                mark.textContent = matchNode.data;
                mark.title = `Highlighted: "${matchNode.data}"`;

                // Replace match with highlight
                matchNode.parentNode.replaceChild(mark, matchNode);

                // Continue with remaining text
                if (afterMatch) {
                    currentNode = afterMatch;
                    regex.lastIndex = 0;
                } else {
                    break;
                }

            } catch (error) {
                console.warn('Error wrapping match:', error);
                break;
            }
        }

        if (iterationCount >= maxIterations) {
            console.warn('Stopped processing due to iteration limit - possible problematic regex pattern');
        }
    }

    // ====================================================================
    // ENHANCED HIGHLIGHTING WITH BETTER ERROR HANDLING
    // ====================================================================

    /**
     * Enhanced highlighting function with better error handling and progress tracking
     * @param {string} query - Search query
     */
    function enhancedHighlightAll(query) {
        if (isHighlighting) {
            console.log('Highlighting already in progress...');
            return;
        }

        const regex = parsePattern(query);
        if (!regex) {
            alert('Invalid search pattern. Use plain text or /regex/flags format.');
            return;
        }

        isHighlighting = true;
        currentPattern = regex;

        try {
            // Clear existing highlights
            clearAllHighlights();

            // Find root element with fallback
            const root = document.body || document.documentElement || document;
            if (!root) {
                throw new Error('No valid root element found');
            }

            // Perform highlighting with progress tracking
            const startTime = performance.now();
            highlightInSubtree(root, regex);
            const endTime = performance.now();

            console.log(`Highlighted pattern "${query}" in ${(endTime - startTime).toFixed(2)}ms`);

            // Setup enhanced persistence
            setupEnhancedPersistence();

        } catch (error) {
            console.error('Error during highlighting:', error);
            alert('Error occurred during highlighting. Check console for details.');
            currentPattern = null;
        } finally {
            isHighlighting = false;
        }
    }

    /**
     * Highlight matches in subtree with enhanced error handling
     * @param {Node} root - Root to search
     * @param {RegExp} regex - Pattern to match
     */
    function highlightInSubtree(root, regex) {
        const walker = createEnhancedTextWalker(root);
        const textNodes = [];

        // Collect text nodes with size limit for performance
        let node;
        let nodeCount = 0;
        const maxNodes = 10000; // Prevent memory issues on huge pages

        while ((node = walker.nextNode()) && nodeCount < maxNodes) {
            textNodes.push(node);
            nodeCount++;
        }

        if (nodeCount >= maxNodes) {
            console.warn('Stopped at node limit for performance. Some text may not be highlighted.');
        }

        // Process nodes in batches for responsiveness
        const batchSize = 50;
        let processedCount = 0;

        function processBatch() {
            const endIndex = Math.min(processedCount + batchSize, textNodes.length);

            for (let i = processedCount; i < endIndex; i++) {
                const textNode = textNodes[i];
                // Revalidate node is still in DOM
                if (textNode.parentNode && textNode.data) {
                    try {
                        enhancedWrapMatches(textNode, regex);
                    } catch (error) {
                        console.warn('Error processing text node:', error);
                    }
                }
            }

            processedCount = endIndex;

            // Schedule next batch if more work remains
            if (processedCount < textNodes.length) {
                scheduleIdleWork(processBatch, { timeout: 100 });
            }
        }

        processBatch();
    }

    // ====================================================================
    // ENHANCED PERSISTENCE WITH THROTTLING
    // ====================================================================

    /**
     * Setup enhanced mutation observer with throttling and coalescing
     */
    function setupEnhancedPersistence() {
        if (!currentPattern) return;

        currentObserver = new MutationObserver((mutations) => {
            // Coalesce mutations to reduce redundant processing
            mutations.forEach(mutation => {
                if (mutation.type === 'childList') {
                    mutation.addedNodes.forEach(node => {
                        if (node.nodeType === Node.TEXT_NODE || node.nodeType === Node.ELEMENT_NODE) {
                            pendingMutations.add(node);
                        }
                    });
                } else if (mutation.type === 'characterData') {
                    pendingMutations.add(mutation.target);
                }
            });

            // Throttle processing to avoid overwhelming the browser
            if (!mutationThrottle && pendingMutations.size > 0) {
                mutationThrottle = scheduleIdleWork(() => {
                    processQueuedMutations();
                    mutationThrottle = null;
                }, { timeout: 200 });
            }
        });

        // Enhanced observation options
        currentObserver.observe(document.body || document.documentElement, {
            childList: true,
            subtree: true,
            characterData: true,
            characterDataOldValue: false // Don't need old values for performance
        });
    }

    /**
     * Process queued mutations efficiently
     */
    function processQueuedMutations() {
        if (!currentPattern || pendingMutations.size === 0) return;

        const mutations = Array.from(pendingMutations);
        pendingMutations.clear();

        mutations.forEach(node => {
            if (!node.parentNode) return; // Node was removed

            try {
                if (node.nodeType === Node.TEXT_NODE) {
                    // Skip if inside existing highlight
                    let parent = node.parentNode;
                    while (parent && parent.nodeType === Node.ELEMENT_NODE) {
                        if (parent.classList && parent.classList.contains(HIGHLIGHT_CLASS)) {
                            return;
                        }
                        parent = parent.parentNode;
                    }
                    enhancedWrapMatches(node, currentPattern);
                } else if (node.nodeType === Node.ELEMENT_NODE) {
                    highlightInSubtree(node, currentPattern);
                }
            } catch (error) {
                console.warn('Error processing mutation:', error);
            }
        });
    }

    // ====================================================================
    // ENHANCED CLEARING WITH BETTER NORMALIZATION
    // ====================================================================

    /**
     * Enhanced clearing with improved text normalization
     */
    function clearAllHighlights() {
        // Disconnect observer during clearing
        if (currentObserver) {
            currentObserver.disconnect();
            currentObserver = null;
        }

        // Clear throttle
        if (mutationThrottle) {
            cancelIdleWork(mutationThrottle);
            mutationThrottle = null;
        }

        pendingMutations.clear();

        try {
            // Enhanced mark removal with batch processing
            const marks = Array.from(document.querySelectorAll(`mark.${HIGHLIGHT_CLASS}`));

            // Process in batches for large numbers of highlights
            const batchSize = 100;
            let processed = 0;

            function clearBatch() {
                const endIndex = Math.min(processed + batchSize, marks.length);

                for (let i = processed; i < endIndex; i++) {
                    const mark = marks[i];
                    if (mark.parentNode) {
                        const textNode = document.createTextNode(mark.textContent);
                        mark.parentNode.replaceChild(textNode, mark);
                    }
                }

                processed = endIndex;

                if (processed < marks.length) {
                    scheduleIdleWork(clearBatch, { timeout: 50 });
                } else {
                    // Final normalization
                    finalizeClearing();
                }
            }

            function finalizeClearing() {
                // Enhanced text normalization
                const root = document.body || document.documentElement;
                if (root && root.normalize) {
                    root.normalize();
                }
                currentPattern = null;
                console.log('All highlights cleared');
            }

            if (marks.length > 0) {
                clearBatch();
            } else {
                finalizeClearing();
            }

        } catch (error) {
            console.error('Error clearing highlights:', error);
            currentPattern = null;
        }
    }

    // ====================================================================
    // ENHANCED KEYBOARD SHORTCUTS
    // ====================================================================

    /**
     * Enhanced keyboard shortcut handling with better key detection
     */
    function handleEnhancedKeyboardShortcuts(event) {
        // Ignore if user is typing in input fields
        const activeElement = document.activeElement;
        const isInputField = activeElement && (
            activeElement.tagName === 'INPUT' ||
            activeElement.tagName === 'TEXTAREA' ||
            activeElement.contentEditable === 'true'
        );

        // Ctrl+Shift+H: New search (case-insensitive)
        if (event.ctrlKey && event.shiftKey && event.key.toLowerCase() === 'h') {
            event.preventDefault();
            event.stopPropagation();

            const currentQuery = currentPattern ? currentPattern.source : '';
            const query = prompt('Enter search pattern (plain text or /regex/flags):', currentQuery);

            if (query !== null && query.trim()) {
                enhancedHighlightAll(query.trim());
            }
        }

        // Ctrl+Shift+C or Escape: Clear highlights (case-insensitive)
        if (((event.ctrlKey && event.shiftKey && event.key.toLowerCase() === 'c') ||
             event.key === 'Escape') && currentPattern) {
            event.preventDefault();
            event.stopPropagation();
            clearAllHighlights();
        }

        // Ctrl+Shift+I: Show current pattern info
        if (event.ctrlKey && event.shiftKey && event.key.toLowerCase() === 'i') {
            event.preventDefault();
            if (currentPattern) {
                const marks = document.querySelectorAll(`mark.${HIGHLIGHT_CLASS}`);
                alert(`Active pattern: ${currentPattern}\nMatches found: ${marks.length}`);
            } else {
                alert('No active highlighting pattern');
            }
        }
    }

    // ====================================================================
    // ENHANCED PUBLIC API
    // ====================================================================

    window.highlightAll = {
        run: enhancedHighlightAll,
        clear: clearAllHighlights,
        getCurrentPattern: () => currentPattern,
        isActive: () => currentPattern !== null,
        getMatchCount: () => document.querySelectorAll(`mark.${HIGHLIGHT_CLASS}`).length,

        // New enhanced methods
        validate: (pattern) => parsePattern(pattern) !== null,
        getStats: () => ({
            pattern: currentPattern ? currentPattern.toString() : null,
            matches: document.querySelectorAll(`mark.${HIGHLIGHT_CLASS}`).length,
            isActive: currentPattern !== null,
            pendingMutations: pendingMutations.size
        })
    };

    // ====================================================================
    // ENHANCED INITIALIZATION
    // ====================================================================

    // Enhanced CSS with better visibility and accessibility
    const enhancedStyle = document.createElement('style');
    enhancedStyle.textContent = `
        .${HIGHLIGHT_CLASS} {
            background-color: #ffff00;
            background-image: linear-gradient(135deg, #ffff00 25%, transparent 25%,
                              transparent 50%, #ffff00 50%, #ffff00 75%,
                              transparent 75%, transparent);
            background-size: 4px 4px;
            color: #000000;
            padding: 1px 2px;
            margin: 0;
            border-radius: 3px;
            border: 1px solid #e6e600;
            box-shadow: 0 1px 2px rgba(0,0,0,0.1);
            position: relative;
            line-height: 1.2;
        }

        .${HIGHLIGHT_CLASS}:hover {
            background-color: #ffff66;
            box-shadow: 0 2px 4px rgba(0,0,0,0.2);
        }

        .${HIGHLIGHT_CLASS}::before {
            content: attr(title);
            position: absolute;
            bottom: 100%;
            left: 50%;
            transform: translateX(-50%);
            background: rgba(0,0,0,0.8);
            color: white;
            padding: 4px 8px;
            border-radius: 4px;
            font-size: 12px;
            white-space: nowrap;
            opacity: 0;
            pointer-events: none;
            transition: opacity 0.2s;
            z-index: 9999;
        }

        .${HIGHLIGHT_CLASS}:hover::before {
            opacity: 1;
        }

        /* High contrast mode support */
        @media (prefers-contrast: high) {
            .${HIGHLIGHT_CLASS} {
                background-color: #000000;
                color: #ffffff;
                border: 2px solid #ffffff;
            }
        }

        /* Reduced motion support */
        @media (prefers-reduced-motion: reduce) {
            .${HIGHLIGHT_CLASS} {
                transition: none;
            }
            .${HIGHLIGHT_CLASS}::before {
                transition: none;
            }
        }
    `;
    document.head.appendChild(enhancedStyle);

    // Enhanced event listeners with better error handling
    document.addEventListener('keydown', handleEnhancedKeyboardShortcuts, true);

    // Cleanup on page unload
    window.addEventListener('beforeunload', () => {
        if (currentObserver) {
            currentObserver.disconnect();
        }
        if (mutationThrottle) {
            cancelIdleWork(mutationThrottle);
        }
    });

    // Enhanced initialization message
    console.log(`
🔍 Highlight-All Enhanced v1.1.0 loaded successfully!

Keyboard shortcuts:
• Ctrl+Shift+H: Search/highlight text
• Ctrl+Shift+C or Escape: Clear highlights
• Ctrl+Shift+I: Show pattern info

API methods:
• window.highlightAll.run(pattern)
• window.highlightAll.clear()
• window.highlightAll.getStats()
• window.highlightAll.validate(pattern)
    `);

})();