ComfyUI-Lora-Manager/static/js/components/ContextMenu/ModelContextMenuMixin.js

import { showToast, getNSFWLevelName, openExampleImagesFolder } from '../../utils/uiHelpers.js';
import { modalManager } from '../../managers/ModalManager.js';
import { state } from '../../state/index.js';
import { getModelApiClient, resetAndReload } from '../../api/modelApiFactory.js';
import { bulkManager } from '../../managers/BulkManager.js';
import { MODEL_CONFIG } from '../../api/apiConfig.js';
import { translate } from '../../utils/i18nHelpers.js';
import { getNsfwLevelSelector } from '../shared/NsfwLevelSelector.js';
import { classifyModelRelinkUrl } from '../../utils/civitaiUtils.js';

// Mixin with shared functionality for LoraContextMenu and CheckpointContextMenu
export const ModelContextMenuMixin = {
    isExcludedView() {
        return state?.pages?.[state.currentPageType]?.viewMode === 'excluded';
    },

    updateExcludeMenuItem() {
        const excludeItem = this.menu?.querySelector('[data-action="exclude"], [data-action="restore"]');
        if (!excludeItem) {
            return;
        }

        const isExcludedView = this.isExcludedView();
        excludeItem.dataset.action = isExcludedView ? 'restore' : 'exclude';
        excludeItem.innerHTML = isExcludedView
            ? `<i class="fas fa-undo"></i> <span>${translate('loras.contextMenu.restoreModel', {}, 'Restore model')}</span>`
            : `<i class="fas fa-eye-slash"></i> <span>${translate('loras.contextMenu.excludeModel', {}, 'Exclude model')}</span>`;
    },

    async restoreExcludedModel(filePath) {
        const restored = await getModelApiClient().unexcludeModel(filePath);
        if (!restored) {
            return;
        }

        if (window.pageControls?.exitExcludedView) {
            await window.pageControls.exitExcludedView();
        } else {
            const resetFn = this.resetAndReload || resetAndReload;
            if (typeof resetFn === 'function') {
                await resetFn(true);
            }
        }
    },

    // NSFW Selector methods
    initNSFWSelector() {
        if (this._nsfwSelectorInitialized) {
            return;
        }

        const selector = getNsfwLevelSelector();
        if (!selector) {
            console.warn('NSFW selector element not found');
            return;
        }

        this._nsfwSelectorInitialized = true;
        this._nsfwSelector = selector;
    },

    resetNSFWSelectorState() {
        // maintained for compatibility; no-op with shared selector
    },

    showNSFWLevelSelector(x, y, card) {
        this.initNSFWSelector();
        const selector = this._nsfwSelector || getNsfwLevelSelector();

        if (!selector) {
            console.warn('NSFW selector not available');
            return;
        }

        // Get current NSFW level
        let currentLevel = 0;
        try {
            const metaData = JSON.parse(card.dataset.meta || '{}');
            currentLevel = metaData.preview_nsfw_level || 0;

            // Update if we have no recorded level but have a dataset attribute
            if (!currentLevel && card.dataset.nsfwLevel) {
                currentLevel = parseInt(card.dataset.nsfwLevel, 10) || 0;
            }
        } catch (err) {
            console.error('Error parsing metadata:', err);
        }

        const filePath = card.dataset.filepath;
        selector.show({
            currentLevel,
            cardPath: filePath,
            onSelect: async (level) => {
                if (!filePath) return false;
                try {
                    await this.saveModelMetadata(filePath, { preview_nsfw_level: level });
                    showToast('toast.contextMenu.contentRatingSet', { level: getNSFWLevelName(level) }, 'success');
                    return true;
                } catch (error) {
                    showToast('toast.contextMenu.contentRatingFailed', { message: error.message }, 'error');
                    return false;
                }
            },
            onClose: () => this.resetNSFWSelectorState(),
        });
    },

    // Civitai re-linking methods
    getModelTypePrefix() {
        // Map the mixin model type to its API route prefix; the relink route
        // exists for all model types via COMMON_ROUTE_DEFINITIONS.
        const prefixMap = {
            lora: 'loras',
            checkpoint: 'checkpoints',
            embedding: 'embeddings'
        };
        return prefixMap[this.modelType] || 'loras';
    },

    showRelinkCivitaiModal() {
        const filePath = this.currentCard.dataset.filepath;
        if (!filePath) return;
        
        // Set up confirm button handler
        const confirmBtn = document.getElementById('confirmRelinkBtn');
        const urlInput = document.getElementById('civitaiModelUrl');
        const errorDiv = document.getElementById('civitaiModelUrlError');
        
        // Remove previous event listener if exists
        if (this._boundRelinkHandler) {
            confirmBtn.removeEventListener('click', this._boundRelinkHandler);
        }
        
        // Create new bound handler
        this._boundRelinkHandler = async () => {
            const url = urlInput.value.trim();
            const { source, modelId, modelVersionId } = classifyModelRelinkUrl(url);

            if (!source || !modelId) {
                errorDiv.textContent = 'Invalid URL format. Expected: https://civitai.com/models/{modelId} or https://civarchive.com/models/{modelId}';
                return;
            }

            errorDiv.textContent = '';
            modalManager.closeModal('relinkCivitaiModal');

            try {
                const isCivArchive = source === 'civarchive';
                state.loadingManager.showSimpleLoading(
                    isCivArchive ? 'Re-linking via CivitArchive...' : 'Re-linking to Civitai...'
                );

                const endpoint = `/api/lm/${this.getModelTypePrefix()}/relink-civitai`;

                const payload = {
                    file_path: filePath,
                    model_id: modelId,
                    model_version_id: modelVersionId
                };
                // Omitted source keeps backend default-provider behaviour; only
                // civarchive pins the provider explicitly.
                if (isCivArchive) {
                    payload.source = source;
                }

                const response = await fetch(endpoint, {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json'
                    },
                    body: JSON.stringify(payload)
                });

                if (!response.ok) {
                    throw new Error(`Failed to re-link model: ${response.statusText}`);
                }

                const data = await response.json();

                if (data.success) {
                    showToast(
                        isCivArchive ? 'toast.contextMenu.linkCivArchSuccess' : 'toast.contextMenu.relinkSuccess',
                        {},
                        'success'
                    );
                    // Reload the current view to show updated data
                    await this.resetAndReload();
                } else {
                    throw new Error(data.error || 'Failed to re-link model');
                }
            } catch (error) {
                console.error('Error re-linking model:', error);
                showToast('toast.contextMenu.relinkFailed', { message: error.message }, 'error');
            } finally {
                state.loadingManager.hide();
            }
        };
        
        // Set new event listener
        confirmBtn.addEventListener('click', this._boundRelinkHandler);
        
        // Clear previous input
        urlInput.value = '';
        errorDiv.textContent = '';
        
        // Show modal
        modalManager.showModal('relinkCivitaiModal');
        
        // Auto-focus the URL input field after modal is shown
        setTimeout(() => urlInput.focus(), 50);
    },

    // HuggingFace linking methods
    showLinkHfModal() {
        const filePath = this.currentCard.dataset.filepath;
        if (!filePath) return;

        const confirmBtn = document.getElementById('confirmLinkHfBtn');
        const urlInput = document.getElementById('hfModelUrl');
        const errorDiv = document.getElementById('hfModelUrlError');

        if (this._boundLinkHfHandler) {
            confirmBtn.removeEventListener('click', this._boundLinkHfHandler);
        }

        this._boundLinkHfHandler = async () => {
            const hfUrl = urlInput.value.trim();
            if (!hfUrl) {
                errorDiv.textContent = 'Please enter a HuggingFace repository URL.';
                return;
            }

            const hfPattern = /^https?:\/\/huggingface\.co\/([^/]+\/[^/]+)\/?$/;
            if (!hfPattern.test(hfUrl)) {
                errorDiv.textContent = 'Invalid URL format. Expected: https://huggingface.co/user/repo';
                return;
            }

            errorDiv.textContent = '';
            modalManager.closeModal('linkHfModal');

            try {
                state.loadingManager.showSimpleLoading('Linking to HuggingFace...');

                const response = await fetch('/api/lm/set-hf-url', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ file_path: filePath, hf_url: hfUrl }),
                });

                if (!response.ok) {
                    const errData = await response.json().catch(() => ({}));
                    throw new Error(errData.error || `Request failed: ${response.statusText}`);
                }

                const data = await response.json();
                if (data.success) {
                    showToast('toast.contextMenu.linkHfSuccess', {}, 'success');
                    await this.resetAndReload();
                } else {
                    throw new Error(data.error || 'Failed to link model');
                }
            } catch (error) {
                console.error('Error linking model to HuggingFace:', error);
                showToast('toast.contextMenu.linkHfFailed', { message: error.message }, 'error');
            } finally {
                state.loadingManager.hide();
            }
        };

        confirmBtn.addEventListener('click', this._boundLinkHfHandler);

        urlInput.value = '';
        errorDiv.textContent = '';

        modalManager.showModal('linkHfModal');

        setTimeout(() => urlInput.focus(), 50);
    },

    // HF metadata enrichment (AI agent) methods
    updateEnrichMenuItem(card) {
        const enrichItem = this.menu?.querySelector('[data-action="enrich-hf-llm"]');
        if (!enrichItem) return;
        const hasHfUrl = !!card.dataset.hf_url;
        enrichItem.classList.toggle('disabled', !hasHfUrl);
        enrichItem.title = hasHfUrl
            ? ''
            : 'Link this model to a HuggingFace repo first (Link Model → Link to HuggingFace)';
    },

    async enrichWithAgent(filePath) {
        const { agentManager } = await import('../../managers/AgentManager.js');

        const configured = await agentManager.isLlmConfigured();
        if (!configured) {
            showToast('toast.agent.llmNotConfigured', {}, 'warning');
            return;
        }

        agentManager.connect();

        const progressUI = state.loadingManager.showEnhancedProgress(
            'Enriching metadata with AI...'
        );

        function cleanupCallbacks() {
            const pIdx = agentManager.progressCallbacks.indexOf(onProgress);
            if (pIdx >= 0) agentManager.progressCallbacks.splice(pIdx, 1);
            const cIdx = agentManager.completeCallbacks.indexOf(onComplete);
            if (cIdx >= 0) agentManager.completeCallbacks.splice(cIdx, 1);
            const eIdx = agentManager.errorCallbacks.indexOf(onError);
            if (eIdx >= 0) agentManager.errorCallbacks.splice(eIdx, 1);
        }

        const onProgress = (data) => {
            if (data.status === 'processing' && data.current_path && data.updated_data && Object.keys(data.updated_data).length > 0) {
                if (state.virtualScroller?.updateSingleItem) {
                    state.virtualScroller.updateSingleItem(data.current_path, data.updated_data);
                }
                const pct = data.total > 0 ? Math.floor((data.processed / data.total) * 100) : 0;
                const name = data.current_path.split('/').pop();
                progressUI.updateProgress(pct, name, `Processing ${name}`);
            }
        };
        agentManager.onProgress(onProgress);

        const onComplete = (data) => {
            cleanupCallbacks();

            if (data.status === 'completed') {
                progressUI.complete(data.summary || 'Enrich complete');
                showToast('toast.agent.enrichComplete', { summary: data.summary || 'Done' }, 'success');
            }
        };
        agentManager.onComplete(onComplete);

        const onError = (data) => {
            cleanupCallbacks();
            state.loadingManager.hide();
            showToast('toast.agent.enrichFailed', { error: data.error || 'Unknown error' }, 'error');
        };
        agentManager.onError(onError);

        try {
            await agentManager.executeSkill('enrich_hf_metadata', [filePath]);
        } catch (error) {
            cleanupCallbacks();
            state.loadingManager.hide();
            showToast('toast.agent.enrichFailed', { error: error.message }, 'error');
        }
    },

    parseModelId(value) {
        if (value === undefined || value === null || value === '') {
            return null;
        }

        const parsed = Number.parseInt(value, 10);
        return Number.isNaN(parsed) ? null : parsed;
    },

    getModelIdFromCard(card) {
        if (!card) {
            return null;
        }

        if (card.dataset?.meta) {
            try {
                const meta = JSON.parse(card.dataset.meta);
                const metaValue = this.parseModelId(meta?.modelId);
                if (metaValue !== null) {
                    return metaValue;
                }
            } catch (error) {
                console.warn('Unable to parse card metadata for model ID', error);
            }
        }

        return null;
    },

    async checkUpdatesForCurrentModel() {
        const card = this.currentCard;
        if (!card) {
            return;
        }

        const modelId = this.getModelIdFromCard(card);
        const typeConfig = MODEL_CONFIG[this.modelType] || {};
        const typeLabel = (typeConfig.displayName || 'Model').toLowerCase();

        if (modelId === null) {
            showToast('toast.models.bulkUpdatesMissing', { type: typeLabel }, 'warning');
            return;
        }

        const apiClient = getModelApiClient();

        const loadingMessage = translate(
            'toast.models.bulkUpdatesChecking',
            { count: 1, type: typeLabel },
            `Checking selected ${typeLabel}(s) for updates...`
        );
        state.loadingManager.showSimpleLoading(loadingMessage);

        try {
            const response = await apiClient.refreshUpdatesForModels([modelId]);
            const records = Array.isArray(response?.records) ? response.records : [];
            const updatesCount = records.length;

            if (updatesCount > 0) {
                showToast('toast.models.bulkUpdatesSuccess', { count: updatesCount, type: typeLabel }, 'success');
            } else {
                showToast('toast.models.bulkUpdatesNone', { type: typeLabel }, 'info');
            }

            const resetFn = this.resetAndReload || resetAndReload;
            if (typeof resetFn === 'function') {
                await resetFn(false);
            }
        } catch (error) {
            console.error('Error checking updates for model:', error);
            showToast(
                'toast.models.bulkUpdatesFailed',
                { type: typeLabel, message: error?.message ?? 'Unknown error' },
                'error'
            );
        } finally {
            state.loadingManager.hide();
            state.loadingManager.restoreProgressBar();
        }
    },

    // Common action handlers
    handleCommonMenuActions(action) {
        switch(action) {
            case 'preview':
                openExampleImagesFolder(this.currentCard.dataset.sha256);
                return true;
            case 'download-examples':
                this.downloadExampleImages(false);
                return true;
            case 'download-examples-force':
                this.downloadExampleImages(true);
                return true;
            case 'civitai':
                if (this.currentCard.dataset.from_civitai === 'true') {
                    if (this.currentCard.querySelector('.fa-globe')) {
                        this.currentCard.querySelector('.fa-globe').click();
                    } else {
                        showToast('toast.contextMenu.fetchMetadataFirst', {}, 'info');
                    }
                } else {
                    showToast('toast.contextMenu.noCivitaiInfo', {}, 'info');
                }
                return true;
            case 'relink-civitai':
                this.showRelinkCivitaiModal();
                return true;
            case 'link-hf':
                this.showLinkHfModal();
                return true;
            case 'enrich-hf-llm':
                this.enrichWithAgent(this.currentCard.dataset.filepath);
                return true;
            case 'set-nsfw':
                this.showNSFWLevelSelector(null, null, this.currentCard);
                return true;
            case 'check-updates':
                this.checkUpdatesForCurrentModel();
                return true;
            default:
                return false;
        }
    },

    // Download example images method
    async downloadExampleImages(force = false) {
        const modelHash = this.currentCard.dataset.sha256;
        if (!modelHash) {
            showToast('toast.contextMenu.missingHash', {}, 'error');
            return;
        }

        try { 
            const apiClient = getModelApiClient();
            await apiClient.downloadExampleImages([modelHash], null, { force });
        } catch (error) {
            console.error('Error downloading example images:', error);
        }
    }
};