(function($){
// === Undo button handler
$('#vmo-undo-btn').off('click').on('click', function () {
        // Try server rollback first if we have a live page; else local undo
        const liveId = (window.currentLivePageId ? parseInt(window.currentLivePageId) : 0) || 0;
        if (!liveId || typeof vmo_ajax === 'undefined') { popUndo(); return; }
        $.post(vmo_ajax.ajaxurl, { action:'vmo_rollback_live_page', security: vmo_ajax.nonce, page_id: liveId, mode: 'check' })
            .done(function(res){
                const count = res?.data?.available || 0;
                if (count < 1) { popUndo(); return; }
                if (!confirm('Undo live change? This will revert the published page to the previous snapshot.')) { return; }
                $.post(vmo_ajax.ajaxurl, { action:'vmo_rollback_live_page', security: vmo_ajax.nonce, page_id: liveId, mode: 'apply' })
                    .done(function(r){
                        if (r.success && r.data && r.data.content) {
                                // Update preview to rolled-back content; keep overlays out of live, but preview uses raw
                                const rolled = r.data.content;
                                window.lastContent = rolled;
                                $('#vmo-code-view').text(rolled);
                                const iframeEl = $('#vmo-preview-iframe')[0];
                                if (iframeEl) {
                                    window.vmoPreviewIsLive = false;
                                    iframeEl.removeAttribute('src');
                                    iframeEl.srcdoc = rolled;
                                    if (typeof window.vmoInjectReplaceButtonsInIframe === 'function') setTimeout(window.vmoInjectReplaceButtonsInIframe, 150);
                                }
                                // Soft refresh builder canvas
                                if (window.vmoState) {
                                        window.vmoState.blocks = [{ id: Date.now(), html: rolled, type: 'rollback' }];
                                        if (typeof renderBlocks === 'function') renderBlocks();
                                }
                                showToast('Live page rolled back ✅', 'success');
                        } else { popUndo(); }
                    })
                    .fail(function(){ popUndo(); });
            })
            .fail(function(){ popUndo(); });
});

// === SAVE ON UNLOAD/NEW PROJECT ===
window.addEventListener('beforeunload', function (e) {
    if (window.isUnsaved) {
        autoSaveCheckpoint();
        e.preventDefault();
        e.returnValue = '';
    }
});

// New Project button
$('#vmo_new_project_btn').off('click').on('click', function () {
    if (window.isUnsaved && !confirm("You have unsaved changes. Start a new project anyway?")) return;

    if (typeof autoSaveCheckpoint === 'function') autoSaveCheckpoint();

    window.lastContent = '';
    window.vmoState = { blocks: [] };
    window.isUnsaved = false;
    window.currentProjectId = null;
    window.currentProjectTitle = null; // <-- RESET the title!
    window.undoStack = [];
    window.vmoChatHistory = [];
    window.vmoImages = [];
    if (window.vmoMediaGenerator && typeof window.vmoMediaGenerator.stopProjectQueuePolling === 'function') {
        window.vmoMediaGenerator.stopProjectQueuePolling();
    }

    // Clear UI
    $('#vmo-code-view').text('');
    if (!window.vmoPreviewIsLive && $('#vmo-preview-iframe').length) $('#vmo-preview-iframe')[0].srcdoc = '';
    $('#builder-canvas').empty();
    $('.vmo-menu-block').remove();
    $('#vmo-menu-modal').hide();
    $('#vmo-menu-items').empty();
    $('#vmo-menu-type').val('Hamburger');
    $('#vmo-menu-position').val('Top Left');
    $('#vmo-menu-logo').val('');
    $('#vmo-menu-extra').val('');
    $('#vmo-prompt-input').val('');
    $('#vmo-load-menu').removeClass('open');

    if (typeof primeUndoStackOnLoad === 'function') primeUndoStackOnLoad(window.vmoState);
    if (typeof renderBlocks === 'function') renderBlocks();

    if (typeof showToast === 'function') showToast('New project started!', 'success');
    console.log('New project state:', {
        lastContent: window.lastContent,
        vmoState: window.vmoState,
        isUnsaved: window.isUnsaved,
        currentProjectId: window.currentProjectId,
        undoStack: window.undoStack,
        vmoChatHistory: window.vmoChatHistory,
        vmoImages: window.vmoImages
    });
});

//=== Save ===
$('#vmo_save_project_btn').off('click').on('click', function () {
    // Check if vmo_ajax is available
    if (typeof vmo_ajax === 'undefined') {
        showToast('System error: AJAX not initialized. Please refresh the page.', 'error');
        return;
    }
    
    let title = window.currentProjectTitle;
    // If new/untitled, prompt for name
    if (!title) {
        title = prompt('Project title:');
        if (!title || !window.lastContent) return showToast('Cannot save empty project.', 'error');
        window.currentProjectTitle = title; // Set global for future updates
    }
    // Always get the latest HTML before saving
    if (window.vmoState && window.vmoState.blocks && window.vmoState.blocks.length) {
        window.lastContent = window.vmoState.blocks.map(b => b.html).join('\n');
    }

    // Enhance for quality: add AOS attrs to headings/images and fix video attributes
    const stripEditorOnlyMediaControls = function(html){
        if (typeof window.vmoStripEditorOnlyMediaOverlays === 'function') {
            return window.vmoStripEditorOnlyMediaOverlays(html);
        }
        return String(html || '').replace(/<button\b[^>]*(?:data-replace-media|data-vmo-editor-only=["']media-replace-overlay["']|class=["'][^"']*(?:vmo-replace-overlay|vmo-replace-btn|vmo-replace-media-btn|replace-image-btn|replace-video-btn)[^"']*["'])[^>]*>[\s\S]*?<\/button>/gi, '');
    };
    const enhanceForQuality = function(html){
        let out = html || '';
        // Add data-aos to headings and images if missing
        out = out.replace(/<(h[1-6]|img)(\b[^>]*?)>/gi, function(m, tag, attrs){
            if (/data-aos=/i.test(attrs)) return m; // already has
            const inject = ' data-aos="fade-up"';
            return `<${tag}${attrs}${inject}>`;
        });
        // Ensure video background best-practices
        out = out.replace(/<video(\b[^>]*?)>/gi, function(m, attrs){
            const needed = ['playsinline','muted','autoplay','loop'];
            let a = attrs;
            needed.forEach(flag=>{ if (!new RegExp('\\b'+flag+'\\b','i').test(a)) a += ' '+flag; });
            return `<video${a}>`;
        });
        return out;
    };
    const enhancedHtml = enhanceForQuality(stripEditorOnlyMediaControls(window.lastContent));
    
    $.post(vmo_ajax.ajaxurl, {
        action: 'vmo_save_project',
        security: vmo_ajax.nonce,
        title,
        html: enhancedHtml,
        state: JSON.stringify(window.vmoState),
        chat: JSON.stringify(window.vmoChatHistory || []),
        images: JSON.stringify(window.vmoImages || [])
    }, function (res) {
        if (res.success) {
            window.isUnsaved = false;
            
            // Update project ID if returned
            if (typeof window.vmoExtractProjectId === 'function') {
                window.currentProjectId = window.vmoExtractProjectId(res) || window.currentProjectId;
            }
            
            // Update preview with local content (not live site)
            try {
                let combinedHtml = '';
                if (window.vmoState && window.vmoState.blocks && window.vmoState.blocks.length) {
                    combinedHtml = window.vmoState.blocks.map(b => stripEditorOnlyMediaControls(b.html)).join('\n');
                } else {
                    combinedHtml = stripEditorOnlyMediaControls(enhancedHtml || '');
                }
                const wrapped = `<div class="vmo-root">${combinedHtml}</div>`;
                
                // Only update local preview if NOT currently viewing live site
                // If viewing live site, show a hint to use "Update Live" button
                const iframeEl = document.getElementById('vmo-preview-iframe');
                if (iframeEl) {
                    if (window.vmoPreviewIsLive && window.currentLivePageId) {
                        // Site is published - don't auto-update live, show hint
                        showToast(`✅ Project "${title}" saved locally.<br><br>Click <strong>Update Live</strong> to push changes to your live website.`, 'success');
                    } else {
                        // Not published yet - update local preview
                        if (typeof window.wrapIframeContent === 'function') {
                            iframeEl.setAttribute('srcdoc', window.wrapIframeContent(wrapped));
                        } else {
                            iframeEl.setAttribute('srcdoc', wrapped);
                        }
                        // Re-inject replace buttons after preview updates
                        iframeEl.addEventListener('load', function once(){
                            iframeEl.removeEventListener('load', once);
                            if (typeof window.vmoInjectReplaceButtonsInIframe === 'function') {
                                setTimeout(window.vmoInjectReplaceButtonsInIframe, 150);
                            }
                        });
                        showToast(`✅ Project "${title}" saved!`, 'success');
                    }
                } else {
                    showToast(`✅ Project "${title}" saved!`, 'success');
                }
            } catch (e) { 
                console.warn('Preview update skipped:', e);
                showToast(`✅ Project "${title}" saved!`, 'success');
            }
        } else {
            showToast(`Failed to save project: ${res.data?.message || 'Unknown error'}`, 'error');
        }
    }).fail(function(xhr, status, error) {
        console.error('Save AJAX Error:', xhr.status, xhr.responseText);
        showToast(`Save failed: ${xhr.status} ${error}. Please refresh the page and try again.`, 'error');
    });
});

// === Load ===
$('#vmo_load_project_btn').off('click').on('click', function (e) {
    e.stopPropagation();
    const $menu = $('#vmo-load-menu');
    if ($menu.hasClass('open')) {
        $menu.removeClass('open');
        return;
    }
    if (window.isGenerating || window.isUnsaved) {
        showToast(`
            ⚠️ You have unsaved changes.<br><br>
            <button id="vmo-confirm-load-btn" class="vmo-toast-action-btn">Load Anyway</button>
            <button id="vmo-save-now-btn" class="vmo-toast-action-btn">💾 Save First</button>
        `, 'error', false, true);
        return;
    }
    loadProjectList();
});

function loadProjectList() {
    const $menu = $('#vmo-load-menu');
    $.post(vmo_ajax.ajaxurl, {
        action: 'vmo_list_projects',
        security: vmo_ajax.nonce
    }, function (res) {
        $menu.empty().addClass('open');
        if (res.success && Array.isArray(res.data.projects) && res.data.projects.length) {
            $menu.data('projects', res.data.projects);
            res.data.projects.forEach((title) => {
                $menu.append(`<button class="vmo-load-btn" data-name="${title}">${title}</button>`);
            });
        } else {
            $menu.append('<div class="vmo-no-proj">No saved projects.</div>');
        }
    });
}

// Load selected project (ALWAYS PARSE STATE AND PRIME UNDO)
$(document).off('click.vmo-load-btn').on('click.vmo-load-btn', '.vmo-load-btn', function () {
    const name = $(this).data('name');
    $.post(vmo_ajax.ajaxurl, {
        action: 'vmo_load_project',
        security: vmo_ajax.nonce,
        name
    }, function (res) {
        if (!res.success) {
            showToast(res.data?.message || 'Failed to load project.', 'error');
            return;
        }
        const proj = res.data;
        let loadedState = proj.state;
        if (typeof loadedState === 'string') {
            try { loadedState = JSON.parse(loadedState); } catch(e){ loadedState = { blocks: [] }; }
        }
        window.vmoState = loadedState && loadedState.blocks ? loadedState : { blocks: [] };
        window.lastContent = proj.html || '';
        window.currentProjectId = proj.project_id || proj.id || null;
        window.currentProjectTitle = name || null; // <-- SET the project title!
        $('#vmo-code-view').text(window.lastContent);
    if (!window.vmoPreviewIsLive) $('#vmo-preview-iframe')[0].srcdoc = window.lastContent;
        primeUndoStackOnLoad(window.vmoState);
        renderBlocks();
        if (window.currentProjectId && window.vmoMediaGenerator && typeof window.vmoMediaGenerator.startProjectQueuePolling === 'function') {
            window.vmoMediaGenerator.startProjectQueuePolling(window.currentProjectId, { source: 'builder' });
        }
        showToast('Project loaded ✅', 'success');
        $('#vmo-load-menu').removeClass('open');
        window.isUnsaved = false;
    });
});

// Handle toast action buttons for save functionality
$(document).off('click.vmo-toast-actions').on('click.vmo-toast-actions', '.vmo-toast-action-btn', function(e) {
    e.stopPropagation();
    const btnId = $(this).attr('id');
    
    if (btnId === 'vmo-save-now-btn') {
        // Save current project
        const title = window.currentProjectTitle || prompt('Project title:');
        if (!title) return;
        
        $.post(vmo_ajax.ajaxurl, {
            action: 'vmo_save_project',
            security: vmo_ajax.nonce,
            title: title,
            html: window.lastContent || '',
            state: JSON.stringify(window.vmoState || { blocks: [] }),
            images: JSON.stringify(window.vmoImages || [])
        }, function (res) {
            if (res.success) {
                showToast('Project saved ✅', 'success');
                window.isUnsaved = false;
                window.currentProjectTitle = title;
                if (typeof window.vmoExtractProjectId === 'function') {
                    window.currentProjectId = window.vmoExtractProjectId(res) || window.currentProjectId;
                }
                // Now proceed with loading
                setTimeout(() => {
                    $('#vmo_load_project_btn').trigger('click');
                }, 500);
            } else {
                showToast('Failed to save project.', 'error');
            }
        });
    } else if (btnId === 'vmo-confirm-load-btn') {
        // Load anyway without saving
        window.isUnsaved = false;
        hideToast(true);
        loadProjectList();
    }
});
})(jQuery);