ComfyUI-Lora-Manager/static/js/managers/BulkMissingLoraDownloadManager.js

import { showToast } from '../utils/uiHelpers.js';
import { isUnresolvableDownloadError } from '../utils/uiHelpers.js';
import { translate } from '../utils/i18nHelpers.js';
import { getModelApiClient } from '../api/modelApiFactory.js';
import { MODEL_TYPES } from '../api/apiConfig.js';
import { extractRecipeId } from '../api/recipeApi.js';
import { state } from '../state/index.js';
import { modalManager } from './ModalManager.js';

/**
 * Manager for downloading missing LoRAs for selected recipes in bulk
 */
export class BulkMissingLoraDownloadManager {
    constructor() {
        this.loraApiClient = getModelApiClient(MODEL_TYPES.LORA);
        this.pendingLoras = [];
        this.pendingRecipes = [];
        this.pendingMissingByRecipe = null;
    }

    /**
     * Collect missing LoRAs from selected recipes with deduplication
     * @param {Array} selectedRecipes - Array of selected recipe objects
     * @returns {Object} - Object containing unique missing LoRAs and statistics
     */
    collectMissingLoras(selectedRecipes) {
        const uniqueLoras = new Map(); // key: hash or modelVersionId, value: lora object
        const missingLorasByRecipe = new Map();
        let totalMissingCount = 0;

        selectedRecipes.forEach(recipe => {
            const missingLoras = [];
            
            if (recipe.loras && Array.isArray(recipe.loras)) {
                recipe.loras.forEach(lora => {
                    // Only include LoRAs not in library and still downloadable
                    // (not deleted from the source, hash still resolvable)
                    if (!lora.inLibrary && !lora.isDeleted && !lora.hashInvalid) {
                        const uniqueKey = lora.hash || lora.id || lora.modelVersionId;
                        
                        if (uniqueKey && !uniqueLoras.has(uniqueKey)) {
                            // Store the LoRA info
                            uniqueLoras.set(uniqueKey, {
                                ...lora,
                                modelId: lora.modelId || lora.model_id,
                                id: lora.id || lora.modelVersionId,
                            });
                        }
                        
                        missingLoras.push(lora);
                        totalMissingCount++;
                    }
                });
            }
            
            if (missingLoras.length > 0) {
                missingLorasByRecipe.set(recipe.id || recipe.file_path, {
                    recipe,
                    missingLoras
                });
            }
        });

        return {
            uniqueLoras: Array.from(uniqueLoras.values()),
            uniqueCount: uniqueLoras.size,
            totalMissingCount,
            missingLorasByRecipe
        };
    }

    /**
     * Show confirmation modal for downloading missing LoRAs
     * @param {Object} stats - Statistics about missing LoRAs
     * @returns {Promise<boolean>} - Whether user confirmed
     */
    async showConfirmationModal(stats) {
        const { uniqueCount, totalMissingCount, uniqueLoras } = stats;

        if (uniqueCount === 0) {
            showToast('toast.recipes.noMissingLoras', {}, 'info');
            return false;
        }

        // Store pending data for confirmation
        this.pendingLoras = uniqueLoras;

        // Update modal content
        const messageEl = document.getElementById('bulkDownloadMissingLorasMessage');
        const listEl = document.getElementById('bulkDownloadMissingLorasList');
        const confirmBtn = document.getElementById('bulkDownloadMissingLorasConfirmBtn');

        if (messageEl) {
            messageEl.textContent = translate('modals.bulkDownloadMissingLoras.message', { 
                uniqueCount, 
                totalCount: totalMissingCount 
            }, `Found ${uniqueCount} unique missing LoRAs (from ${totalMissingCount} total across selected recipes).`);
        }

        if (listEl) {
            listEl.innerHTML = uniqueLoras.slice(0, 10).map(lora => `
                <li>
                    <span class="lora-name">${lora.name || lora.file_name || 'Unknown'}</span>
                    ${lora.version ? `<span class="lora-version">${lora.version}</span>` : ''}
                </li>
            `).join('') + 
            (uniqueLoras.length > 10 ? `
                <li class="more-items">${translate('modals.bulkDownloadMissingLoras.moreItems', { count: uniqueLoras.length - 10 }, `...and ${uniqueLoras.length - 10} more`)}</li>
            ` : '');
        }

        if (confirmBtn) {
            confirmBtn.innerHTML = `
                <i class="fas fa-download"></i>
                ${translate('modals.bulkDownloadMissingLoras.downloadButton', { count: uniqueCount }, `Download ${uniqueCount} LoRA(s)`)}
            `;
        }

        // Show modal
        modalManager.showModal('bulkDownloadMissingLorasModal');
        
        // Return a promise that will be resolved when user confirms or cancels
        return new Promise((resolve) => {
            this.confirmResolve = resolve;
        });
    }

    /**
     * Called when user confirms download in modal
     */
    async confirmDownload() {
        modalManager.closeModal('bulkDownloadMissingLorasModal');
        
        if (this.confirmResolve) {
            this.confirmResolve(true);
            this.confirmResolve = null;
        }

        // Execute download
        await this.executeDownload(this.pendingLoras);
        this.pendingLoras = [];
        this.pendingMissingByRecipe = null;
    }

    /**
     * Download missing LoRAs for selected recipes
     * @param {Array} selectedRecipes - Array of selected recipe objects
     */
    async downloadMissingLoras(selectedRecipes) {
        if (!selectedRecipes || selectedRecipes.length === 0) {
            showToast('toast.recipes.noRecipesSelected', {}, 'warning');
            return;
        }

        // Store selected recipes
        this.pendingRecipes = selectedRecipes;

        // Collect missing LoRAs with deduplication
        const stats = this.collectMissingLoras(selectedRecipes);
        // Kept so executeDownload can mark unresolvable failures back onto
        // every recipe occurrence (hashInvalid → reconnect candidacy).
        this.pendingMissingByRecipe = stats.missingLorasByRecipe;
        
        if (stats.uniqueCount === 0) {
            showToast('toast.recipes.noMissingLorasInSelection', {}, 'info');
            return;
        }

        // Show confirmation modal
        const confirmed = await this.showConfirmationModal(stats);
        if (!confirmed) {
            return;
        }
    }

    /**
     * Execute the download process
     * @param {Array} lorasToDownload - Array of unique LoRAs to download
     */
    async executeDownload(lorasToDownload) {
        const totalLoras = lorasToDownload.length;
        
        // Get LoRA root directory
        const loraRoot = await this.getLoraRoot();
        if (!loraRoot) {
            showToast('toast.recipes.noLoraRootConfigured', {}, 'error');
            return;
        }

        // Generate batch download ID
        const batchDownloadId = Date.now().toString();
        
        // Use default paths
        const useDefaultPaths = true;

        // Set up WebSocket for progress updates
        const wsProtocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://';
        const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/download-progress?id=${batchDownloadId}`);

        // Show download progress UI
        const loadingManager = state.loadingManager;
        const updateProgress = loadingManager.showDownloadProgress(totalLoras);

        let completedDownloads = 0;
        let failedDownloads = 0;
        let markedInvalidCount = 0;
        let currentLoraProgress = 0;
        let cancelled = false;

        loadingManager.showCancelButton(async () => {
            if (cancelled) return;
            cancelled = true;
            try {
                await this.loraApiClient.cancelDownload(batchDownloadId);
            } catch (e) {
                console.error('Cancel request failed:', e);
            }
        });

        // Set up WebSocket message handler
        ws.onmessage = (event) => {
            const data = JSON.parse(event.data);

            // Handle download ID confirmation
            if (data.type === 'download_id') {
                console.log(`Connected to batch download progress with ID: ${data.download_id}`);
                return;
            }

            if (data.status === 'cancelled') {
                cancelled = true;
                return;
            }

            // Process progress updates
            if (data.status === 'progress' && data.download_id && data.download_id.startsWith(batchDownloadId)) {
                currentLoraProgress = data.progress;
                
                const currentLora = lorasToDownload[completedDownloads + failedDownloads];
                const loraName = currentLora ? (currentLora.name || currentLora.file_name || 'Unknown') : '';

                const metrics = {
                    bytesDownloaded: data.bytes_downloaded,
                    totalBytes: data.total_bytes,
                    bytesPerSecond: data.bytes_per_second
                };

                updateProgress(currentLoraProgress, completedDownloads, loraName, metrics);

                // Update status message
                if (currentLoraProgress < 3) {
                    loadingManager.setStatus(
                        translate('recipes.controls.import.startingDownload', 
                            { current: completedDownloads + failedDownloads + 1, total: totalLoras },
                            `Starting download for LoRA ${completedDownloads + failedDownloads + 1}/${totalLoras}`
                        )
                    );
                } else if (currentLoraProgress > 3 && currentLoraProgress < 100) {
                    loadingManager.setStatus(
                        translate('recipes.controls.import.downloadingLoras', {}, `Downloading LoRAs...`)
                    );
                }
            }
        };

        // Wait for WebSocket to connect
        await new Promise((resolve, reject) => {
            ws.onopen = resolve;
            ws.onerror = (error) => {
                console.error('WebSocket error:', error);
                reject(error);
            };
        });

        // Download each LoRA sequentially
        for (let i = 0; i < lorasToDownload.length; i++) {
            if (cancelled) break;

            const lora = lorasToDownload[i];
            
            currentLoraProgress = 0;
            
            loadingManager.setStatus(
                translate('recipes.controls.import.startingDownload', 
                    { current: i + 1, total: totalLoras },
                    `Starting download for LoRA ${i + 1}/${totalLoras}`
                )
            );
            updateProgress(0, completedDownloads, lora.name || lora.file_name || 'Unknown');

            try {
                const modelId = lora.modelId || lora.model_id;
                const versionId = lora.id || lora.modelVersionId;

                if (!modelId && !versionId) {
                    console.warn(`Skipping LoRA without model/version ID:`, lora);
                    failedDownloads++;
                    continue;
                }

                const response = await this.loraApiClient.downloadModel(
                    modelId,
                    versionId,
                    loraRoot,
                    '',
                    useDefaultPaths,
                    batchDownloadId
                );

                if (cancelled) break;

                if (!response.success) {
                    console.error(`Failed to download LoRA ${lora.name || lora.file_name}: ${response.error}`);
                    failedDownloads++;
                    // An unresolvable failure (model gone on CivitAI) flips
                    // every recipe occurrence to reconnect candidacy — same
                    // rule as the single-LoRA download in RecipeModal.
                    if (isUnresolvableDownloadError(response.error)) {
                        markedInvalidCount += await this.markLoraHashInvalidInRecipes(lora);
                    }
                } else {
                    completedDownloads++;
                    updateProgress(100, completedDownloads, '');
                }
            } catch (error) {
                if (!cancelled) {
                    console.error(`Error downloading LoRA ${lora.name || lora.file_name}:`, error);
                    failedDownloads++;
                    if (isUnresolvableDownloadError(error?.message)) {
                        markedInvalidCount += await this.markLoraHashInvalidInRecipes(lora);
                    }
                }
            }
        }

        // Close WebSocket
        ws.close();

        // Hide loading UI
        loadingManager.hide();

        // Show completion message
        if (cancelled) {
            showToast('toast.downloads.downloadStopped', {}, 'info',
                `Download cancelled. ${completedDownloads} item(s) completed.`);
        } else if (failedDownloads === 0) {
            showToast('toast.loras.allDownloadSuccessful', { count: completedDownloads }, 'success');
        } else {
            showToast('toast.loras.downloadPartialSuccess', {
                completed: completedDownloads,
                total: totalLoras
            }, 'warning');
        }

        // Unresolvable failures were marked hash-invalid during the loop;
        // tell the user those entries now offer reconnect instead of download.
        if (markedInvalidCount > 0) {
            showToast('toast.recipes.unresolvableMarkedForReconnect', {
                count: markedInvalidCount
            }, 'info', `${markedInvalidCount} unresolvable entr(ies) marked — they can now be reconnected to a local LoRA.`);
        }

        // Update each affected recipe card with fresh data (LoRA inLibrary flags changed)
        if (state.virtualScroller) {
            for (const recipe of this.pendingRecipes) {
                const recipeId = extractRecipeId(recipe.file_path);
                if (!recipeId) continue;
                try {
                    const detailRes = await fetch(`/api/lm/recipe/${encodeURIComponent(recipeId)}`);
                    if (detailRes.ok) {
                        const updated = await detailRes.json();
                        state.virtualScroller.updateSingleItem(recipe.file_path, updated);
                    }
                } catch (e) {
                    console.warn('Failed to update recipe card after LoRA download:', e);
                }
            }
        }
    }

    /**
     * Mark every recipe occurrence of a failed LoRA as hash-invalid.
     *
     * Mirrors RecipeModal.markLoraHashInvalid for the bulk flow: the flag
     * makes each occurrence an unresolved rematch candidate and swaps its
     * action from download to reconnect. Only called for unresolvable
     * failures — transient errors leave entries untouched.
     *
     * @param {Object} failedLora - The deduplicated LoRA that failed
     * @returns {Promise<number>} - How many recipe entries were marked
     */
    async markLoraHashInvalidInRecipes(failedLora) {
        const failedKey = failedLora.hash || failedLora.id || failedLora.modelVersionId;
        if (!failedKey || !this.pendingMissingByRecipe) {
            return 0;
        }

        let marked = 0;
        for (const { recipe, missingLoras } of this.pendingMissingByRecipe.values()) {
            const recipeId = extractRecipeId(recipe.file_path) || recipe.id;
            if (!recipeId || !Array.isArray(recipe.loras)) {
                continue;
            }
            for (const entry of missingLoras) {
                const entryKey = entry.hash || entry.id || entry.modelVersionId;
                if (entryKey !== failedKey) {
                    continue;
                }
                const loraIndex = recipe.loras.indexOf(entry);
                if (loraIndex < 0) {
                    continue;
                }
                try {
                    const response = await fetch('/api/lm/recipe/lora/mark-hash-invalid', {
                        method: 'POST',
                        headers: { 'Content-Type': 'application/json' },
                        body: JSON.stringify({
                            recipe_id: recipeId,
                            lora_index: loraIndex,
                        }),
                    });
                    if (response.ok) {
                        entry.hashInvalid = true;
                        marked++;
                    }
                } catch (error) {
                    console.warn('Failed to mark LoRA hash invalid:', error);
                }
            }
        }
        return marked;
    }

    /**
     * Get LoRA root directory from API
     * @returns {Promise<string|null>} - LoRA root directory or null
     */
    async getLoraRoot() {
        try {
            // Fetch available LoRA roots from API
            const rootsData = await this.loraApiClient.fetchModelRoots();
            
            if (!rootsData || !rootsData.roots || rootsData.roots.length === 0) {
                console.error('No LoRA roots available');
                return null;
            }

            // Try to get default root from settings
            const defaultRootKey = 'default_lora_root';
            const defaultRoot = state.global?.settings?.[defaultRootKey];
            
            // If default root is set and exists in available roots, use it
            if (defaultRoot && rootsData.roots.includes(defaultRoot)) {
                return defaultRoot;
            }
            
            // Otherwise, return the first available root
            return rootsData.roots[0];
            
        } catch (error) {
            console.error('Error getting LoRA root:', error);
            return null;
        }
    }
}

// Export singleton instance
export const bulkMissingLoraDownloadManager = new BulkMissingLoraDownloadManager();

// Make available globally for HTML onclick handlers
if (typeof window !== 'undefined') {
    window.bulkMissingLoraDownloadManager = bulkMissingLoraDownloadManager;
}