static/js/main.js

// Multi-Modal Knowledge Distillation - JavaScript

class KnowledgeDistillationApp {
    constructor() {
        this.selectedModels = [];
        this.selectedTeachers = [];
        this.selectedStudent = null;
        this.configuredModels = {};
        this.currentStep = 1;
        this.trainingSession = null;
        this.websocket = null;

        // Add global error handler
        window.addEventListener('error', (event) => {
            console.error('Global error:', event.error);
            this.handleGlobalError(event.error);
        });

        // Add unhandled promise rejection handler
        window.addEventListener('unhandledrejection', (event) => {
            console.error('Unhandled promise rejection:', event.reason);
            this.handleGlobalError(event.reason);
        });

        this.init();
    }

    handleGlobalError(error) {
        const errorMsg = error?.message || 'An unexpected error occurred';
        console.error('Handling global error:', errorMsg);

        // Try to show error in UI, fallback to console
        try {
            if (this.showError) {
                this.showError(`Error: ${errorMsg}`);
            }
        } catch (e) {
            console.error('Could not show error in UI:', e);
        }
    }
    
    init() {
        this.setupEventListeners();
        this.updateModelCount();

        // Initialize models manager
        this.modelsManager = new ModelsManager(this);
    }
    
    setupEventListeners() {
        // File upload
        const uploadArea = document.getElementById('upload-area');
        const fileInput = document.getElementById('file-input');
        
        uploadArea.addEventListener('click', () => fileInput.click());
        uploadArea.addEventListener('dragover', this.handleDragOver.bind(this));
        uploadArea.addEventListener('dragleave', this.handleDragLeave.bind(this));
        uploadArea.addEventListener('drop', this.handleDrop.bind(this));
        fileInput.addEventListener('change', this.handleFileSelect.bind(this));
        
        // Hugging Face models
        document.getElementById('add-hf-model').addEventListener('click', this.addHuggingFaceModel.bind(this));
        document.getElementById('hf-repo').addEventListener('keypress', (e) => {
            if (e.key === 'Enter') this.addHuggingFaceModel();
        });
        
        // URL models
        document.getElementById('add-url-model').addEventListener('click', this.addUrlModel.bind(this));
        document.getElementById('model-url').addEventListener('keypress', (e) => {
            if (e.key === 'Enter') this.addUrlModel();
        });
        
        // Navigation
        document.getElementById('next-step-1').addEventListener('click', () => this.goToStep(2));
        document.getElementById('back-step-2').addEventListener('click', () => this.goToStep(1));
        document.getElementById('back-step-3').addEventListener('click', () => this.goToStep(2));
        document.getElementById('start-training').addEventListener('click', this.showConfirmModal.bind(this));
        document.getElementById('start-new-training').addEventListener('click', () => this.resetAndGoToStep(1));
        
        // Training controls
        document.getElementById('cancel-training').addEventListener('click', this.cancelTraining.bind(this));
        document.getElementById('download-model').addEventListener('click', this.downloadModel.bind(this));
        
        // Modals
        document.getElementById('confirm-start').addEventListener('click', this.startTraining.bind(this));
        document.getElementById('confirm-cancel').addEventListener('click', this.hideConfirmModal.bind(this));
        document.getElementById('error-ok').addEventListener('click', this.hideErrorModal.bind(this));

        // Suggested models
        document.querySelectorAll('.suggestion-btn').forEach(btn => {
            btn.addEventListener('click', (e) => {
                const modelName = e.target.getAttribute('data-model');
                const trustRequired = e.target.classList.contains('trust-required');
                const gatedModel = e.target.classList.contains('gated-model');

                document.getElementById('hf-repo').value = modelName;

                // Auto-enable trust remote code if required
                if (trustRequired) {
                    document.getElementById('trust-remote-code').checked = true;
                    this.showTokenStatus('⚠️ Trust Remote Code enabled for this model', 'warning');
                }

                // Show warning for gated models
                if (gatedModel) {
                    const tokenInput = document.getElementById('hf-token');
                    if (!tokenInput.value.trim()) {
                        this.showTokenStatus('🔒 This model requires a Hugging Face token and access permission!', 'error');
                        tokenInput.focus();
                        return;
                    } else {
                        this.showTokenStatus('✅ Token detected for gated model', 'success');
                    }
                }

                this.addHuggingFaceModel();
            });
        });

        // Test token button
        document.getElementById('test-token').addEventListener('click', this.testToken.bind(this));

        // Test model button
        document.getElementById('test-model').addEventListener('click', this.testModel.bind(this));

        // Download and upload buttons
        document.getElementById('download-model').addEventListener('click', this.downloadModel.bind(this));
        document.getElementById('upload-to-hf').addEventListener('click', this.showHFUploadModal.bind(this));
        document.getElementById('confirm-hf-upload').addEventListener('click', this.uploadToHuggingFace.bind(this));
        document.getElementById('cancel-hf-upload').addEventListener('click', this.hideHFUploadModal.bind(this));

        // Incremental training
        document.getElementById('enable-incremental').addEventListener('change', this.toggleIncrementalTraining.bind(this));
        document.getElementById('existing-student').addEventListener('change', this.onStudentModelChange.bind(this));
        document.getElementById('refresh-students').addEventListener('click', this.loadTrainedStudents.bind(this));

        // Student source options
        document.querySelectorAll('input[name="student-source"]').forEach(radio => {
            radio.addEventListener('change', this.onStudentSourceChange.bind(this));
        });

        // HF student model
        document.getElementById('test-student-model').addEventListener('click', this.testStudentModel.bind(this));
        document.getElementById('add-hf-student').addEventListener('click', this.addHFStudentModel.bind(this));

        // HF Space student model
        document.getElementById('test-space-model').addEventListener('click', this.testSpaceModel.bind(this));
        document.getElementById('add-space-student').addEventListener('click', this.addSpaceStudentModel.bind(this));

        // File upload
        document.getElementById('student-file-upload').addEventListener('change', this.onStudentFilesUpload.bind(this));

        // Load trained students on page load
        this.loadTrainedStudents();
    }
    
    // File handling
    handleDragOver(e) {
        e.preventDefault();
        e.currentTarget.classList.add('dragover');
    }
    
    handleDragLeave(e) {
        e.preventDefault();
        e.currentTarget.classList.remove('dragover');
    }
    
    handleDrop(e) {
        e.preventDefault();
        e.currentTarget.classList.remove('dragover');
        const files = Array.from(e.dataTransfer.files);
        this.processFiles(files);
    }
    
    handleFileSelect(e) {
        const files = Array.from(e.target.files);
        this.processFiles(files);
    }
    
    async processFiles(files) {
        const validFiles = files.filter(file => this.validateFile(file));

        if (validFiles.length === 0) {
            this.showError('No valid model files selected. Please select .pt, .pth, .bin, or .safetensors files.');
            return;
        }

        this.showLoading(`Processing ${validFiles.length} file(s)...`);

        try {
            for (const file of validFiles) {
                await this.uploadFile(file);
            }
        } catch (error) {
            this.showError(`Error processing files: ${error.message}`);
        } finally {
            this.hideLoading();
        }
    }
    
    validateFile(file) {
        const validExtensions = ['.pt', '.pth', '.bin', '.safetensors'];
        const extension = '.' + file.name.split('.').pop().toLowerCase();
        const maxSize = 5 * 1024 * 1024 * 1024; // 5GB
        
        if (!validExtensions.includes(extension)) {
            this.showError(`Invalid file type: ${file.name}. Allowed types: ${validExtensions.join(', ')}`);
            return false;
        }
        
        if (file.size > maxSize) {
            this.showError(`File too large: ${file.name}. Maximum size: 5GB`);
            return false;
        }
        
        return true;
    }
    
    async uploadFile(file) {
        const formData = new FormData();
        formData.append('files', file);
        formData.append('model_names', file.name.split('.')[0]);

        try {
            const response = await fetch('/upload', {
                method: 'POST',
                body: formData
            });

            if (!response.ok) {
                throw new Error(`HTTP error! status: ${response.status}`);
            }

            const result = await response.json();

            if (result.success) {
                result.models.forEach(model => this.addModel(model));
                this.addConsoleMessage(`Successfully uploaded: ${file.name}`, 'success');
            } else {
                throw new Error(result.message || 'Upload failed');
            }
        } catch (error) {
            this.showError(`Upload failed for ${file.name}: ${error.message}`);
            throw error;
        }
    }
    
    async addHuggingFaceModel() {
        const repoInput = document.getElementById('hf-repo');
        const tokenInput = document.getElementById('hf-token');
        const accessTypeSelect = document.getElementById('model-access-type');

        const repo = repoInput.value.trim();
        const manualToken = tokenInput.value.trim();
        const accessType = accessTypeSelect ? accessTypeSelect.value : 'read';

        if (!repo) {
            this.showError('Please enter a Hugging Face repository name');
            return;
        }

        if (!this.isValidHuggingFaceRepo(repo)) {
            this.showError('Invalid repository format. Use format: organization/model-name (e.g., google/bert_uncased_L-2_H-128_A-2)');
            return;
        }

        let tokenToUse = manualToken;

        // If no manual token provided, get appropriate token for access type
        if (!manualToken) {
            try {
                const response = await fetch(`/api/tokens/for-task/${accessType}`);
                if (response.ok) {
                    const data = await response.json();
                    if (data.success) {
                        // We don't store the actual token, just indicate it will be used
                        this.showSuccess(`سيتم استخدام ${data.token_info.type_name} للوصول للنموذج`);
                        tokenToUse = 'auto'; // Indicate automatic token selection
                    }
                } else {
                    this.showWarning('لم يتم العثور على رمز مناسب، قد تحتاج لإضافة رمز يدوياً');
                }
            } catch (error) {
                console.error('Error getting token for task:', error);
                this.showWarning('خطأ في الحصول على الرمز المناسب');
            }
        }

        const model = {
            id: `hf_${Date.now()}`,
            name: repo,
            source: 'huggingface',
            path: repo,
            token: tokenToUse,
            accessType: accessType,
            info: { modality: 'unknown', format: 'huggingface' }
        };

        this.addModel(model);
        repoInput.value = '';
        // Don't clear token as user might want to use it for multiple models
    }
    
    async addUrlModel() {
        const urlInput = document.getElementById('model-url');
        const url = urlInput.value.trim();

        if (!url) {
            this.showError('Please enter a model URL');
            return;
        }

        if (!this.isValidUrl(url)) {
            this.showError('Invalid URL format');
            return;
        }

        // Validate that URL points to a model file
        const filename = this.extractFilenameFromUrl(url);
        const validExtensions = ['.pt', '.pth', '.bin', '.safetensors'];
        const hasValidExtension = validExtensions.some(ext => filename.toLowerCase().endsWith(ext));

        if (!hasValidExtension) {
            this.showError(`URL must point to a model file with extension: ${validExtensions.join(', ')}`);
            return;
        }

        this.showLoading('Validating URL...');

        try {
            // Test if URL is accessible
            const response = await fetch(url, { method: 'HEAD' });
            if (!response.ok) {
                throw new Error(`URL not accessible: ${response.status}`);
            }

            const model = {
                id: `url_${Date.now()}`,
                name: filename,
                source: 'url',
                path: url,
                info: {
                    modality: 'unknown',
                    format: filename.split('.').pop(),
                    size: response.headers.get('content-length') ? parseInt(response.headers.get('content-length')) : null
                }
            };

            this.addModel(model);
            urlInput.value = '';
            this.hideLoading();

        } catch (error) {
            this.hideLoading();
            this.showError(`URL validation failed: ${error.message}`);
        }
    }
    
    addModel(model) {
        if (this.selectedModels.length >= 10) {
            this.showError('Maximum 10 models allowed');
            return;
        }
        
        // Check for duplicates
        if (this.selectedModels.some(m => m.path === model.path)) {
            this.showError('Model already added');
            return;
        }
        
        this.selectedModels.push(model);
        this.updateModelsDisplay();
        this.updateModelCount();
        this.updateNextButton();
    }
    
    removeModel(modelId) {
        this.selectedModels = this.selectedModels.filter(m => m.id !== modelId);
        this.updateModelsDisplay();
        this.updateModelCount();
        this.updateNextButton();
    }
    
    updateModelsDisplay() {
        const grid = document.getElementById('models-grid');
        grid.innerHTML = '';
        
        this.selectedModels.forEach(model => {
            const card = this.createModelCard(model);
            grid.appendChild(card);
        });
    }
    
    createModelCard(model) {
        const card = document.createElement('div');
        card.className = 'model-card';
        
        const modalityIcon = this.getModalityIcon(model.info.modality);
        const sizeText = model.size ? this.formatBytes(model.size) : 'Unknown size';
        
        card.innerHTML = `
            <button class="model-remove" onclick="app.removeModel('${model.id}')">×</button>
            <h4>${modalityIcon} ${model.name}</h4>
            <div class="model-info">Source: ${model.source}</div>
            <div class="model-info">Format: ${model.info.format}</div>
            <div class="model-info">Modality: ${model.info.modality}</div>
            <div class="model-info">Size: ${sizeText}</div>
        `;
        
        return card;
    }
    
    getModalityIcon(modality) {
        const icons = {
            text: '<i class="fas fa-font"></i>',
            vision: '<i class="fas fa-eye"></i>',
            multimodal: '<i class="fas fa-layer-group"></i>',
            audio: '<i class="fas fa-volume-up"></i>',
            unknown: '<i class="fas fa-question"></i>'
        };
        return icons[modality] || icons.unknown;
    }
    
    updateModelCount() {
        document.getElementById('model-count').textContent = this.selectedModels.length;
    }
    
    updateNextButton() {
        const button = document.getElementById('next-step-1');
        button.disabled = this.selectedModels.length === 0;
    }
    
    // Navigation
    goToStep(step) {
        // Hide all steps
        document.querySelectorAll('.step-section').forEach(section => {
            section.classList.add('hidden');
        });

        // Show target step
        document.getElementById(`step-${step}`).classList.remove('hidden');
        this.currentStep = step;
    }

    resetAndGoToStep(step) {
        // Reset training session
        this.trainingSession = null;
        if (this.websocket) {
            this.websocket.close();
            this.websocket = null;
        }

        // Reset UI elements
        document.getElementById('download-model').classList.add('hidden');
        document.getElementById('start-new-training').classList.add('hidden');
        document.getElementById('cancel-training').classList.remove('hidden');

        // Clear console
        document.getElementById('training-console').innerHTML = '';

        // Reset progress
        document.getElementById('overall-progress').style.width = '0%';
        document.getElementById('progress-percentage').textContent = '0%';

        // Go to step
        this.goToStep(step);
    }
    
    // Training
    showConfirmModal() {
        document.getElementById('confirm-modal').classList.remove('hidden');
    }
    
    hideConfirmModal() {
        document.getElementById('confirm-modal').classList.add('hidden');
    }
    
    async startTraining() {
        this.hideConfirmModal();

        // Get configuration
        const config = this.getTrainingConfig();

        // Check if any models require token and warn user
        const hasGatedModels = this.selectedModels.some(model =>
            model.path.includes('gemma') ||
            model.path.includes('llama') ||
            model.path.includes('claude')
        );

        if (hasGatedModels && !config.hf_token) {
            const proceed = confirm(
                'Some selected models may require a Hugging Face token for access. ' +
                'Do you want to continue without a token? (Training may fail for gated models)'
            );
            if (!proceed) return;
        }

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

            const result = await response.json();

            if (result.success) {
                this.trainingSession = result.session_id;
                this.goToStep(3);
                this.connectWebSocket();
                this.startProgressPolling();
            } else {
                throw new Error(result.message || 'Failed to start training');
            }
        } catch (error) {
            this.showError(`Failed to start training: ${error.message}`);
        }
    }
    
    getTrainingConfig() {
        // Get HF token from interface
        const hfToken = document.getElementById('hf-token').value.trim();
        const trustRemoteCode = document.getElementById('trust-remote-code').checked;
        const incrementalTraining = document.getElementById('enable-incremental').checked;
        const existingStudent = document.getElementById('existing-student').value;

        // Get student model info based on source
        let studentModelPath = null;
        let studentSource = 'local';

        if (incrementalTraining && existingStudent) {
            const selectedOption = document.querySelector('#existing-student option:checked');
            if (selectedOption && selectedOption.dataset.source === 'huggingface') {
                studentSource = 'huggingface';
                studentModelPath = existingStudent; // Already the repo name
            } else if (selectedOption && selectedOption.dataset.source === 'space') {
                studentSource = 'space';
                studentModelPath = existingStudent.startsWith('space:') ? existingStudent.substring(6) : existingStudent;
            } else {
                studentSource = 'local';
                studentModelPath = existingStudent;
            }
        }

        const config = {
            session_id: `session_${Date.now()}`,
            teacher_models: this.selectedModels.map(m => ({
                path: m.path,
                token: m.token || hfToken || null,
                trust_remote_code: trustRemoteCode
            })),
            student_config: {
                hidden_size: parseInt(document.getElementById('hidden-size').value),
                num_layers: parseInt(document.getElementById('num-layers').value),
                output_size: parseInt(document.getElementById('hidden-size').value)
            },
            training_params: {
                max_steps: parseInt(document.getElementById('max-steps').value),
                learning_rate: parseFloat(document.getElementById('learning-rate').value),
                temperature: parseFloat(document.getElementById('temperature').value),
                alpha: parseFloat(document.getElementById('alpha').value),
                batch_size: 8
            },
            distillation_strategy: document.getElementById('strategy').value,
            hf_token: hfToken || null,
            trust_remote_code: trustRemoteCode,
            incremental_training: incrementalTraining,
            existing_student_model: studentModelPath,
            student_source: studentSource
        };

        return config;
    }
    
    connectWebSocket() {
        if (!this.trainingSession) return;
        
        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
        const wsUrl = `${protocol}//${window.location.host}/ws/${this.trainingSession}`;
        
        this.websocket = new WebSocket(wsUrl);
        
        this.websocket.onmessage = (event) => {
            const data = JSON.parse(event.data);
            if (data.type === 'training_update') {
                this.updateTrainingProgress(data.data);
            }
        };
        
        this.websocket.onerror = (error) => {
            console.error('WebSocket error:', error);
            this.addConsoleMessage('WebSocket connection error', 'error');
        };
        
        this.websocket.onclose = () => {
            console.log('WebSocket connection closed');
        };
    }
    
    async startProgressPolling() {
        if (!this.trainingSession) return;

        this.trainingStartTime = Date.now(); // Track start time

        const poll = async () => {
            try {
                const response = await fetch(`/progress/${this.trainingSession}`);
                const progress = await response.json();
                this.updateTrainingProgress(progress);

                // If stuck on loading for too long, show helpful message
                if (progress.status === 'loading_models' && progress.progress < 0.2) {
                    const elapsed = Date.now() - this.trainingStartTime;
                    if (elapsed > 60000) { // 1 minute
                        const messageEl = document.getElementById('training-message');
                        if (messageEl && !messageEl.innerHTML.includes('Large models')) {
                            messageEl.innerHTML = `${progress.message}<br><small style="color: #666;">Large models may take several minutes to load. Please be patient...</small>`;
                        }
                    }
                }

                if (progress.status === 'completed' || progress.status === 'failed') {
                    return; // Stop polling
                }

                setTimeout(poll, 2000); // Poll every 2 seconds
            } catch (error) {
                console.error('Error polling progress:', error);
                setTimeout(poll, 5000); // Retry after 5 seconds
            }
        };

        poll();
    }
    
    updateTrainingProgress(progress) {
        // Update progress bar
        const progressFill = document.getElementById('overall-progress');
        const progressText = document.getElementById('progress-percentage');
        const percentage = Math.round(progress.progress * 100);
        
        progressFill.style.width = `${percentage}%`;
        progressText.textContent = `${percentage}%`;
        
        // Update status info
        document.getElementById('training-status').textContent = this.formatStatus(progress.status);
        document.getElementById('current-step').textContent = `${progress.current_step} / ${progress.total_steps}`;
        document.getElementById('eta').textContent = progress.eta || 'Calculating...';
        
        // Update metrics
        if (progress.loss !== null && progress.loss !== undefined) {
            document.getElementById('current-loss').textContent = progress.loss.toFixed(4);
        }
        
        // Add console message
        if (progress.message) {
            this.addConsoleMessage(progress.message, this.getMessageType(progress.status));
        }
        
        // Handle completion
        if (progress.status === 'completed') {
            document.getElementById('download-model').classList.remove('hidden');
            document.getElementById('upload-to-hf').classList.remove('hidden');
            document.getElementById('start-new-training').classList.remove('hidden');
            document.getElementById('cancel-training').classList.add('hidden');
            this.addConsoleMessage('Training completed successfully!', 'success');
        } else if (progress.status === 'failed') {
            document.getElementById('start-new-training').classList.remove('hidden');
            document.getElementById('cancel-training').classList.add('hidden');
            this.addConsoleMessage(`Training failed: ${progress.message}`, 'error');
        }
    }
    
    formatStatus(status) {
        const statusMap = {
            'initializing': 'Initializing...',
            'loading_models': 'Loading Models...',
            'initializing_student': 'Initializing Student...',
            'training': 'Training...',
            'saving': 'Saving Model...',
            'completed': 'Completed',
            'failed': 'Failed'
        };
        return statusMap[status] || status;
    }
    
    getMessageType(status) {
        if (status === 'completed') return 'success';
        if (status === 'failed') return 'error';
        if (status === 'loading_models' || status === 'initializing') return 'warning';
        return 'info';
    }
    
    addConsoleMessage(message, type = 'info') {
        const console = document.getElementById('training-console');
        if (!console) {
            // Fallback to browser console if training console not found
            console.log(`[${type.toUpperCase()}] ${message}`);
            return;
        }

        try {
            const line = document.createElement('div');
            line.className = `console-line ${type}`;
            line.textContent = `[${new Date().toLocaleTimeString()}] ${message}`;
            console.appendChild(line);
            console.scrollTop = console.scrollHeight;
        } catch (error) {
            console.error('Error adding console message:', error);
            console.log(`[${type.toUpperCase()}] ${message}`);
        }
    }
    
    async cancelTraining() {
        if (this.websocket) {
            this.websocket.close();
        }
        this.addConsoleMessage('Training cancelled by user', 'warning');
    }
    
    async downloadModel() {
        if (!this.trainingSession) return;
        
        try {
            const response = await fetch(`/download/${this.trainingSession}`);
            if (response.ok) {
                const blob = await response.blob();
                const url = window.URL.createObjectURL(blob);
                const a = document.createElement('a');
                a.href = url;
                a.download = `distilled_model_${this.trainingSession}.safetensors`;
                document.body.appendChild(a);
                a.click();
                document.body.removeChild(a);
                window.URL.revokeObjectURL(url);
            } else {
                throw new Error('Download failed');
            }
        } catch (error) {
            this.showError(`Download failed: ${error.message}`);
        }
    }
    
    // Utility functions
    isValidHuggingFaceRepo(repo) {
        return /^[a-zA-Z0-9_.-]+\/[a-zA-Z0-9_.-]+$/.test(repo);
    }
    
    isValidUrl(url) {
        try {
            new URL(url);
            return true;
        } catch {
            return false;
        }
    }
    
    extractFilenameFromUrl(url) {
        try {
            const pathname = new URL(url).pathname;
            return pathname.split('/').pop() || 'model';
        } catch {
            return 'model';
        }
    }
    
    formatBytes(bytes) {
        const sizes = ['B', 'KB', 'MB', 'GB', 'TB'];
        if (bytes === 0) return '0 B';
        const i = Math.floor(Math.log(bytes) / Math.log(1024));
        return `${(bytes / Math.pow(1024, i)).toFixed(1)} ${sizes[i]}`;
    }
    
    showError(message) {
        try {
            const errorMessage = document.getElementById('error-message');
            const errorModal = document.getElementById('error-modal');

            if (errorMessage && errorModal) {
                errorMessage.textContent = message;
                errorModal.classList.remove('hidden');
            } else {
                // Fallback: use alert if modal elements not found
                console.error('Error modal elements not found, using alert');
                alert(`Error: ${message}`);
            }
        } catch (error) {
            console.error('Error showing error message:', error);
            alert(`Error: ${message}`);
        }
    }
    
    hideErrorModal() {
        document.getElementById('error-modal').classList.add('hidden');
    }
    
    showLoading(message) {
        // Create loading overlay if it doesn't exist
        let loadingOverlay = document.getElementById('loading-overlay');
        if (!loadingOverlay) {
            loadingOverlay = document.createElement('div');
            loadingOverlay.id = 'loading-overlay';
            loadingOverlay.className = 'loading-overlay';
            loadingOverlay.innerHTML = `
                <div class="loading-content">
                    <div class="loading-spinner"></div>
                    <div class="loading-message">${message}</div>
                </div>
            `;
            document.body.appendChild(loadingOverlay);
        } else {
            loadingOverlay.querySelector('.loading-message').textContent = message;
            loadingOverlay.classList.remove('hidden');
        }
    }

    hideLoading() {
        const loadingOverlay = document.getElementById('loading-overlay');
        if (loadingOverlay) {
            loadingOverlay.classList.add('hidden');
        }
    }

    async testToken() {
        const tokenInput = document.getElementById('hf-token');
        const statusDiv = document.getElementById('token-status');
        const token = tokenInput.value.trim();

        if (!token) {
            this.showTokenStatus('Please enter a token first', 'warning');
            return;
        }

        this.showLoading('Testing token...');

        try {
            const response = await fetch('/test-token');
            const result = await response.json();

            this.hideLoading();

            if (result.token_valid) {
                this.showTokenStatus('✅ Token is valid and working!', 'success');
            } else if (result.token_available) {
                this.showTokenStatus(`❌ Token validation failed: ${result.message}`, 'error');
            } else {
                this.showTokenStatus('⚠️ No token found in environment. Using interface token.', 'warning');
            }
        } catch (error) {
            this.hideLoading();
            this.showTokenStatus(`❌ Error testing token: ${error.message}`, 'error');
        }
    }

    showTokenStatus(message, type) {
        const statusDiv = document.getElementById('token-status');
        if (!statusDiv) {
            console.warn('Token status div not found, using console message instead');
            console.log(`${type.toUpperCase()}: ${message}`);
            return;
        }

        statusDiv.textContent = message;
        statusDiv.className = `token-status ${type}`;
        statusDiv.classList.remove('hidden');

        // Hide after 5 seconds
        setTimeout(() => {
            if (statusDiv) {
                statusDiv.classList.add('hidden');
            }
        }, 5000);
    }

    async testModel() {
        const repoInput = document.getElementById('hf-repo');
        const trustRemoteCode = document.getElementById('trust-remote-code').checked;
        const repo = repoInput.value.trim();

        if (!repo) {
            this.showTokenStatus('Please enter a model repository name first', 'warning');
            return;
        }

        if (!this.isValidHuggingFaceRepo(repo)) {
            this.showTokenStatus('Invalid repository format. Use: organization/model-name', 'error');
            return;
        }

        this.showLoading(`Testing model: ${repo}...`);

        try {
            const response = await fetch('/test-model', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    model_path: repo,
                    trust_remote_code: trustRemoteCode
                })
            });

            const result = await response.json();
            this.hideLoading();

            if (result.success) {
                const info = result.model_info;
                let message = `✅ Model ${repo} is accessible!`;
                if (info.architecture) {
                    message += ` Architecture: ${info.architecture}`;
                }
                if (info.modality) {
                    message += `, Modality: ${info.modality}`;
                }
                this.showTokenStatus(message, 'success');
            } else {
                let message = `❌ Model test failed: ${result.error}`;
                if (result.suggestions && result.suggestions.length > 0) {
                    message += `. Suggestions: ${result.suggestions.join(', ')}`;
                }
                this.showTokenStatus(message, 'error');
            }
        } catch (error) {
            this.hideLoading();
            this.showTokenStatus(`❌ Error testing model: ${error.message}`, 'error');
        }
    }

    downloadModel() {
        if (!this.trainingSession) {
            this.showError('No training session found');
            return;
        }

        // Create download link
        const downloadUrl = `/download/${this.trainingSession}`;
        const link = document.createElement('a');
        link.href = downloadUrl;
        link.download = `distilled_model_${this.trainingSession}`;
        document.body.appendChild(link);
        link.click();
        document.body.removeChild(link);

        this.addConsoleMessage('Download started...', 'info');
    }

    showHFUploadModal() {
        const modal = document.getElementById('hf-upload-modal');
        modal.classList.remove('hidden');

        // Pre-fill token if available
        const hfToken = document.getElementById('hf-token').value.trim();
        if (hfToken) {
            document.getElementById('hf-upload-token').value = hfToken;
            // Auto-validate token and suggest username
            this.validateTokenAndSuggestName(hfToken);
        }
    }

    hideHFUploadModal() {
        const modal = document.getElementById('hf-upload-modal');
        modal.classList.add('hidden');
    }

    async uploadToHuggingFace() {
        if (!this.trainingSession) {
            this.showError('No training session found');
            return;
        }

        const repoName = document.getElementById('hf-repo-name').value.trim();
        const description = document.getElementById('hf-description').value.trim();
        const token = document.getElementById('hf-upload-token').value.trim();
        const isPrivate = document.getElementById('hf-private').checked;

        if (!repoName || !token) {
            this.showError('Repository name and token are required');
            return;
        }

        if (!repoName.includes('/')) {
            this.showError('Repository name must be in format: username/model-name');
            return;
        }

        this.showLoading('Uploading model to Hugging Face...');
        this.hideHFUploadModal();

        try {
            const formData = new FormData();
            formData.append('repo_name', repoName);
            formData.append('description', description);
            formData.append('private', isPrivate);
            formData.append('hf_token', token);

            const response = await fetch(`/upload-to-hf/${this.trainingSession}`, {
                method: 'POST',
                body: formData
            });

            const result = await response.json();
            this.hideLoading();

            if (result.success) {
                this.addConsoleMessage(`✅ Model uploaded successfully to ${result.repo_url}`, 'success');
                this.addConsoleMessage(`📁 Uploaded files: ${result.uploaded_files.join(', ')}`, 'info');

                // Show success message with link
                const successMsg = document.createElement('div');
                successMsg.className = 'alert alert-success';
                successMsg.innerHTML = `
                    <strong>🎉 Upload Successful!</strong><br>
                    Your model is now available at: <a href="${result.repo_url}" target="_blank">${result.repo_url}</a>
                `;

                // Find a safe container to insert the message
                let container = document.querySelector('.step-3 .step-content');
                if (!container) {
                    container = document.querySelector('.step-3');
                }
                if (!container) {
                    container = document.querySelector('#training-progress');
                }
                if (!container) {
                    container = document.body;
                }

                if (container && container.firstChild) {
                    container.insertBefore(successMsg, container.firstChild);
                } else if (container) {
                    container.appendChild(successMsg);
                }

                // Remove after 10 seconds
                setTimeout(() => {
                    if (successMsg && successMsg.parentNode) {
                        successMsg.parentNode.removeChild(successMsg);
                    }
                }, 10000);

            } else {
                const errorMsg = result.detail || result.message || 'Unknown error';
                this.showError(`Upload failed: ${errorMsg}`);
                this.addConsoleMessage(`❌ Upload failed: ${errorMsg}`, 'error');
            }

        } catch (error) {
            this.hideLoading();
            const errorMsg = error.message || 'Network error occurred';
            this.showError(`Upload failed: ${errorMsg}`);
            this.addConsoleMessage(`❌ Upload error: ${errorMsg}`, 'error');
            console.error('Upload error details:', error);
        }
    }

    async loadTrainedStudents() {
        try {
            const response = await fetch('/trained-students');
            const data = await response.json();

            const select = document.getElementById('existing-student');
            select.innerHTML = '<option value="">Select a trained model...</option>';

            if (data.trained_students && data.trained_students.length > 0) {
                data.trained_students.forEach(model => {
                    const option = document.createElement('option');
                    option.value = model.path;
                    option.textContent = `${model.name} (${model.architecture}, ${model.training_sessions} sessions)`;
                    option.dataset.modelInfo = JSON.stringify(model);
                    select.appendChild(option);
                });
            } else {
                const option = document.createElement('option');
                option.value = '';
                option.textContent = 'No trained models found';
                option.disabled = true;
                select.appendChild(option);
            }
        } catch (error) {
            console.error('Error loading trained students:', error);
            const select = document.getElementById('existing-student');
            select.innerHTML = '<option value="">Error loading models</option>';
        }
    }

    toggleIncrementalTraining() {
        const enabled = document.getElementById('enable-incremental').checked;
        const options = document.getElementById('incremental-options');

        if (enabled) {
            options.classList.remove('hidden');
            this.loadTrainedStudents();
        } else {
            options.classList.add('hidden');
            document.getElementById('student-info').classList.add('hidden');
        }
    }

    onStudentModelChange() {
        const select = document.getElementById('existing-student');
        const selectedOption = select.options[select.selectedIndex];
        const studentInfo = document.getElementById('student-info');

        if (selectedOption && selectedOption.dataset.modelInfo) {
            const modelData = JSON.parse(selectedOption.dataset.modelInfo);

            // Update info display
            document.getElementById('student-arch').textContent = modelData.architecture || 'Unknown';
            document.getElementById('student-teachers').textContent =
                modelData.original_teachers.length > 0 ?
                modelData.original_teachers.join(', ') :
                'None';
            document.getElementById('student-sessions').textContent = modelData.training_sessions || '0';
            document.getElementById('student-last').textContent =
                modelData.last_training !== 'unknown' ?
                new Date(modelData.last_training).toLocaleString() :
                'Unknown';

            studentInfo.classList.remove('hidden');
        } else {
            studentInfo.classList.add('hidden');
        }
    }

    onStudentSourceChange() {
        try {
            const selectedRadio = document.querySelector('input[name="student-source"]:checked');
            if (!selectedRadio) {
                console.warn('No student source radio button selected');
                return;
            }

            const selectedSource = selectedRadio.value;

            // Hide all options safely
            const optionIds = ['local-student-options', 'hf-student-options', 'space-student-options', 'upload-student-options'];
            optionIds.forEach(id => {
                const element = document.getElementById(id);
                if (element) {
                    element.classList.add('hidden');
                }
            });

            // Show selected option
            const targetElement = document.getElementById(`${selectedSource}-student-options`);
            if (targetElement) {
                targetElement.classList.remove('hidden');
            } else {
                console.warn(`Element ${selectedSource}-student-options not found`);
            }

            // Reset student info
            const studentInfo = document.getElementById('student-info');
            if (studentInfo) {
                studentInfo.classList.add('hidden');
            }
        } catch (error) {
            console.error('Error in onStudentSourceChange:', error);
        }
    }

    async testStudentModel() {
        const repoInput = document.getElementById('hf-student-repo');
        const repo = repoInput.value.trim();

        if (!repo) {
            this.showTokenStatus('Please enter a student model repository name', 'warning');
            return;
        }

        if (!this.isValidHuggingFaceRepo(repo)) {
            this.showTokenStatus('Invalid repository format. Use: organization/model-name', 'error');
            return;
        }

        this.showLoading(`Testing student model: ${repo}...`);

        try {
            const response = await fetch('/test-model', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    model_path: repo,
                    trust_remote_code: document.getElementById('trust-remote-code').checked
                })
            });

            const result = await response.json();
            this.hideLoading();

            if (result.success) {
                this.showTokenStatus(`✅ Student model ${repo} is accessible!`, 'success');
            } else {
                this.showTokenStatus(`❌ Student model test failed: ${result.error}`, 'error');
            }
        } catch (error) {
            this.hideLoading();
            this.showTokenStatus(`❌ Error testing student model: ${error.message}`, 'error');
        }
    }

    addHFStudentModel() {
        const repo = document.getElementById('hf-student-repo').value.trim();

        if (!repo) {
            this.showTokenStatus('Please enter a repository name first', 'warning');
            return;
        }

        if (!this.isValidHuggingFaceRepo(repo)) {
            this.showTokenStatus('Invalid repository format. Use: organization/model-name', 'error');
            return;
        }

        // Set the HF repo as the selected student model
        const existingStudentSelect = document.getElementById('existing-student');

        // Remove any existing HF options to avoid duplicates
        Array.from(existingStudentSelect.options).forEach(option => {
            if (option.value.startsWith('hf:')) {
                option.remove();
            }
        });

        // Add HF repo as an option
        const option = document.createElement('option');
        option.value = repo; // Store the repo directly, not with hf: prefix
        option.textContent = `${repo} (Hugging Face)`;
        option.selected = true;
        option.dataset.source = 'huggingface';
        existingStudentSelect.appendChild(option);

        // Update student info display
        this.displayHFStudentInfo(repo);

        // Show success message
        this.showTokenStatus(`✅ Added Hugging Face student model: ${repo}`, 'success');

        // Clear input
        document.getElementById('hf-student-repo').value = '';
    }

    displayHFStudentInfo(repo) {
        // Show student info for HF model
        const studentInfo = document.getElementById('student-info');

        document.getElementById('student-arch').textContent = 'Hugging Face Model';
        document.getElementById('student-teachers').textContent = 'Unknown (External Model)';
        document.getElementById('student-sessions').textContent = 'N/A';
        document.getElementById('student-last').textContent = 'External Model';

        studentInfo.classList.remove('hidden');

        // Add note about HF model
        const noteDiv = document.createElement('div');
        noteDiv.className = 'alert alert-info';
        noteDiv.innerHTML = `
            <i class="fas fa-info-circle"></i>
            <strong>Hugging Face Model:</strong> ${repo}<br>
            This model will be loaded from Hugging Face Hub. Make sure you have access to it.
        `;

        // Remove any existing notes
        const existingNotes = studentInfo.querySelectorAll('.alert-info');
        existingNotes.forEach(note => note.remove());

        studentInfo.appendChild(noteDiv);
    }

    async testSpaceModel() {
        const spaceInput = document.getElementById('hf-space-repo');
        const space = spaceInput.value.trim();

        if (!space) {
            this.showTokenStatus('Please enter a Space name first', 'warning');
            return;
        }

        if (!this.isValidHuggingFaceRepo(space)) {
            this.showTokenStatus('Invalid Space format. Use: username/space-name', 'error');
            return;
        }

        this.showLoading(`Testing Space: ${space}...`);

        try {
            // Test if the Space exists and has models
            const response = await fetch('/test-space', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    space_name: space,
                    hf_token: document.getElementById('hf-token').value.trim()
                })
            });

            const result = await response.json();
            this.hideLoading();

            if (result.success) {
                const modelsCount = result.models ? result.models.length : 0;
                this.showTokenStatus(`✅ Space ${space} is accessible! Found ${modelsCount} trained models.`, 'success');
            } else {
                this.showTokenStatus(`❌ Space test failed: ${result.error}`, 'error');
            }
        } catch (error) {
            this.hideLoading();
            this.showTokenStatus(`❌ Error testing Space: ${error.message}`, 'error');
        }
    }

    addSpaceStudentModel() {
        const space = document.getElementById('hf-space-repo').value.trim();

        if (!space) {
            this.showTokenStatus('Please enter a Space name first', 'warning');
            return;
        }

        if (!this.isValidHuggingFaceRepo(space)) {
            this.showTokenStatus('Invalid Space format. Use: username/space-name', 'error');
            return;
        }

        // Set the Space as the selected student model
        const existingStudentSelect = document.getElementById('existing-student');

        // Remove any existing Space options to avoid duplicates
        Array.from(existingStudentSelect.options).forEach(option => {
            if (option.value.startsWith('space:')) {
                option.remove();
            }
        });

        // Add Space as an option
        const option = document.createElement('option');
        option.value = `space:${space}`;
        option.textContent = `${space} (Hugging Face Space)`;
        option.selected = true;
        option.dataset.source = 'space';
        existingStudentSelect.appendChild(option);

        // Update student info display
        this.displaySpaceStudentInfo(space);

        // Show success message
        this.showTokenStatus(`✅ Added Hugging Face Space: ${space}`, 'success');

        // Clear input
        document.getElementById('hf-space-repo').value = '';
    }

    displaySpaceStudentInfo(space) {
        // Show student info for Space
        const studentInfo = document.getElementById('student-info');

        document.getElementById('student-arch').textContent = 'Hugging Face Space';
        document.getElementById('student-teachers').textContent = 'Multiple Models Available';
        document.getElementById('student-sessions').textContent = 'External Space';
        document.getElementById('student-last').textContent = 'External Space';

        studentInfo.classList.remove('hidden');

        // Add note about Space
        const noteDiv = document.createElement('div');
        noteDiv.className = 'alert alert-info';
        noteDiv.innerHTML = `
            <i class="fas fa-rocket"></i>
            <strong>Hugging Face Space:</strong> ${space}<br>
            This will load trained models from another Space. The Space should have completed training and saved models.
        `;

        // Remove any existing notes
        const existingNotes = studentInfo.querySelectorAll('.alert-info');
        existingNotes.forEach(note => note.remove());

        studentInfo.appendChild(noteDiv);
    }

    onStudentFilesUpload(event) {
        const files = event.target.files;
        if (files.length === 0) return;

        const fileNames = Array.from(files).map(f => f.name);
        this.showTokenStatus(`📁 Selected files: ${fileNames.join(', ')}`, 'success');

        // TODO: Implement file upload functionality
        // For now, just show that files were selected
    }

    async validateTokenAndSuggestName(token) {
        if (!token) return;

        try {
            const response = await fetch('/validate-repo-name', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    repo_name: 'test/test', // Dummy name to get username
                    hf_token: token
                })
            });

            const result = await response.json();

            if (result.username) {
                // Auto-suggest repository name
                const repoInput = document.getElementById('hf-repo-name');
                if (!repoInput.value.trim()) {
                    const modelName = `distilled-model-${Date.now()}`;
                    repoInput.value = `${result.username}/${modelName}`;
                    repoInput.placeholder = `${result.username}/your-model-name`;
                }
            }
        } catch (error) {
            console.error('Error validating token:', error);
        }
    }

    async validateRepoName() {
        const repoName = document.getElementById('hf-repo-name').value.trim();
        const token = document.getElementById('hf-upload-token').value.trim();

        if (!repoName || !token) return;

        try {
            const response = await fetch('/validate-repo-name', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    repo_name: repoName,
                    hf_token: token
                })
            });

            const result = await response.json();

            const statusDiv = document.getElementById('repo-validation-status');
            if (!statusDiv) {
                // Create status div if it doesn't exist
                const div = document.createElement('div');
                div.id = 'repo-validation-status';
                div.className = 'validation-status';
                document.getElementById('hf-repo-name').parentNode.appendChild(div);
            }

            const status = document.getElementById('repo-validation-status');

            if (result.valid) {
                status.innerHTML = `✅ Repository name is valid`;
                status.className = 'validation-status success';
            } else {
                status.innerHTML = `❌ ${result.error}`;
                if (result.suggested_name) {
                    status.innerHTML += `<br>💡 Suggested: <strong>${result.suggested_name}</strong>`;
                    // Auto-fill suggested name
                    document.getElementById('hf-repo-name').value = result.suggested_name;
                }
                status.className = 'validation-status error';
            }

            status.classList.remove('hidden');

        } catch (error) {
            console.error('Error validating repo name:', error);
        }
    }
}

// Initialize app when DOM is loaded
document.addEventListener('DOMContentLoaded', () => {
    window.app = new KnowledgeDistillationApp();
});

// Advanced Features Functions
async function showGoogleModels() {
    try {
        const response = await fetch('/api/models/google');
        const data = await response.json();

        if (response.ok) {
            const modelsHtml = data.models.map(model => `
                <div class="model-card">
                    <h4>${model.name}</h4>
                    <p>${model.description}</p>
                    <div class="model-info">
                        <span class="badge ${model.medical_specialized ? 'bg-success' : 'bg-info'}">
                            ${model.medical_specialized ? 'Medical Specialized' : 'General Purpose'}
                        </span>
                        <span class="badge bg-secondary">${model.size_gb} GB</span>
                        <span class="badge bg-primary">${model.modality}</span>
                    </div>
                    <button class="btn btn-primary mt-2" onclick="addGoogleModel('${model.name}')">
                        Add to Teachers
                    </button>
                </div>
            `).join('');

            showModal('Google Models', modelsHtml);
        }
    } catch (error) {
        console.error('Error loading Google models:', error);
        showError('Failed to load Google models');
    }
}

async function showSystemInfo() {
    try {
        const response = await fetch('/api/system/performance');
        const data = await response.json();

        if (response.ok) {
            const systemInfoHtml = `
                <div class="system-info">
                    <h5>Memory Information</h5>
                    <div class="info-grid">
                        <div class="info-item">
                            <strong>Process Memory:</strong> ${data.memory.process_memory_mb.toFixed(1)} MB
                        </div>
                        <div class="info-item">
                            <strong>Memory Usage:</strong> ${data.memory.process_memory_percent.toFixed(1)}%
                        </div>
                        <div class="info-item">
                            <strong>Available Memory:</strong> ${data.memory.system_memory_available_gb.toFixed(1)} GB
                        </div>
                        <div class="info-item">
                            <strong>CPU Cores:</strong> ${data.cpu_cores}
                        </div>
                    </div>

                    <h5 class="mt-3">Optimizations Applied</h5>
                    <ul class="optimization-list">
                        ${data.optimizations_applied.map(opt => `<li>${opt}</li>`).join('')}
                    </ul>

                    ${data.recommendations.length > 0 ? `
                        <h5 class="mt-3">Recommendations</h5>
                        <ul class="recommendation-list">
                            ${data.recommendations.map(rec => `<li>${rec}</li>`).join('')}
                        </ul>
                    ` : ''}

                    <div class="mt-3">
                        <button class="btn btn-warning" onclick="forceMemoryCleanup()">
                            Force Memory Cleanup
                        </button>
                    </div>
                </div>
            `;

            showModal('System Information', systemInfoHtml);
        }
    } catch (error) {
        console.error('Error loading system info:', error);
        showError('Failed to load system information');
    }
}

async function forceMemoryCleanup() {
    try {
        const response = await fetch('/api/system/cleanup', { method: 'POST' });
        const data = await response.json();

        if (response.ok) {
            showSuccess(data.message);
            // Refresh system info
            setTimeout(() => showSystemInfo(), 1000);
        } else {
            showError('Failed to cleanup memory');
        }
    } catch (error) {
        console.error('Error during memory cleanup:', error);
        showError('Error during memory cleanup');
    }
}

function addGoogleModel(modelName) {
    // Add the Google model to the HF repo input
    const hfRepoInput = document.getElementById('hf-repo');
    if (hfRepoInput) {
        hfRepoInput.value = modelName;
        // Trigger the add model function
        if (window.app && window.app.addHuggingFaceModel) {
            window.app.addHuggingFaceModel();
        }
    }
    closeModal();
}

function showModal(title, content) {
    // Create modal if it doesn't exist
    let modal = document.getElementById('advanced-modal');
    if (!modal) {
        modal = document.createElement('div');
        modal.id = 'advanced-modal';
        modal.className = 'modal-overlay';
        modal.innerHTML = `
            <div class="modal-content">
                <div class="modal-header">
                    <h3 id="modal-title">${title}</h3>
                    <button class="modal-close" onclick="closeModal()">&times;</button>
                </div>
                <div class="modal-body" id="modal-body">
                    ${content}
                </div>
            </div>
        `;
        document.body.appendChild(modal);
    } else {
        document.getElementById('modal-title').textContent = title;
        document.getElementById('modal-body').innerHTML = content;
    }

    modal.style.display = 'flex';
}

function closeModal() {
    const modal = document.getElementById('advanced-modal');
    if (modal) {
        modal.style.display = 'none';
    }
}

function showSuccess(message) {
    showNotification(message, 'success');
}

function showError(message) {
    showNotification(message, 'error');
}

function showNotification(message, type) {
    const notification = document.createElement('div');
    notification.className = `notification notification-${type}`;
    notification.textContent = message;

    document.body.appendChild(notification);

    // Auto remove after 5 seconds
    setTimeout(() => {
        if (notification.parentNode) {
            notification.parentNode.removeChild(notification);
        }
    }, 5000);
}

// Models Management Functions
class ModelsManager {
    constructor(app) {
        this.app = app;
        this.setupEventListeners();
        this.loadConfiguredModels();
    }

    setupEventListeners() {
        // Refresh models
        const refreshButton = document.getElementById('refresh-models');
        if (refreshButton) {
            refreshButton.addEventListener('click', () => {
                this.loadConfiguredModels();
            });
        }

        // Search models
        const searchButton = document.getElementById('search-models-btn');
        if (searchButton) {
            searchButton.addEventListener('click', () => {
                this.searchModels();
            });
        }

        // Search on Enter key
        const searchQuery = document.getElementById('model-search-query');
        if (searchQuery) {
            searchQuery.addEventListener('keypress', (e) => {
                if (e.key === 'Enter') {
                    this.searchModels();
                }
            });
        }

        // Add custom model
        const addCustomButton = document.getElementById('add-custom-model');
        if (addCustomButton) {
            addCustomButton.addEventListener('click', () => {
                this.showAddCustomModelModal();
            });
        }
    }

    async loadConfiguredModels() {
        try {
            const response = await fetch('/api/models/teachers');
            const data = await response.json();

            if (data.success) {
                this.app.configuredModels = data.teachers;
                this.app.selectedTeachers = data.selected;
                this.displayConfiguredModels(data.teachers, data.selected);
            }

        } catch (error) {
            console.error('Error loading configured models:', error);
            this.app.showError('خطأ في تحميل النماذج المُعدة');
        }
    }

    displayConfiguredModels(models, selected) {
        const container = document.getElementById('configured-models-list');

        if (!container) return;

        if (Object.keys(models).length === 0) {
            container.innerHTML = '<p class="text-muted">لا توجد نماذج مُعدة</p>';
            return;
        }

        container.innerHTML = Object.entries(models).map(([id, model]) => `
            <div class="card mb-2">
                <div class="card-body">
                    <div class="d-flex justify-content-between align-items-start">
                        <div class="flex-grow-1">
                            <div class="form-check">
                                <input class="form-check-input" type="checkbox"
                                       id="model-${id}" ${selected.includes(id) ? 'checked' : ''}
                                       onchange="window.app.modelsManager.toggleModelSelection('${id}', this.checked)">
                                <label class="form-check-label" for="model-${id}">
                                    <h6 class="mb-1">${model.name}</h6>
                                </label>
                            </div>
                            <p class="text-muted small mb-1">${model.description || 'لا يوجد وصف'}</p>
                            <div class="d-flex gap-2">
                                <span class="badge bg-primary">${model.category}</span>
                                <span class="badge bg-secondary">${model.modality}</span>
                                <span class="badge bg-info">${model.parameters || 'Unknown'}</span>
                            </div>
                        </div>
                        <div class="d-flex gap-1">
                            <button class="btn btn-sm btn-outline-info" onclick="window.app.modelsManager.showModelInfo('${id}')">
                                <i class="fas fa-info"></i>
                            </button>
                            <button class="btn btn-sm btn-outline-danger" onclick="window.app.modelsManager.removeModel('${id}')">
                                <i class="fas fa-trash"></i>
                            </button>
                        </div>
                    </div>
                </div>
            </div>
        `).join('');
    }

    async searchModels() {
        const queryElement = document.getElementById('model-search-query');
        const typeElement = document.getElementById('model-type-filter');

        if (!queryElement) return;

        const query = queryElement.value.trim();
        const modelType = typeElement ? typeElement.value : '';

        if (!query) {
            this.app.showError('يرجى إدخال كلمة البحث');
            return;
        }

        const searchButton = document.getElementById('search-models-btn');
        const originalText = searchButton.innerHTML;
        searchButton.innerHTML = '<i class="fas fa-spinner fa-spin"></i> جاري البحث...';
        searchButton.disabled = true;

        try {
            const response = await fetch('/api/models/search', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({
                    query: query,
                    limit: 20,
                    model_type: modelType || null
                })
            });

            const data = await response.json();

            if (data.success) {
                this.displaySearchResults(data.results);
            } else {
                this.app.showError('فشل في البحث عن النماذج');
            }

        } catch (error) {
            console.error('Error searching models:', error);
            this.app.showError('خطأ في البحث عن النماذج');
        } finally {
            searchButton.innerHTML = originalText;
            searchButton.disabled = false;
        }
    }

    displaySearchResults(results) {
        const resultsContainer = document.getElementById('model-search-results-list');
        const searchResults = document.getElementById('model-search-results');

        if (!resultsContainer || !searchResults) return;

        if (results.length === 0) {
            resultsContainer.innerHTML = '<p class="text-muted">لم يتم العثور على نتائج</p>';
        } else {
            resultsContainer.innerHTML = results.map(result => `
                <div class="card mb-2">
                    <div class="card-body">
                        <div class="d-flex justify-content-between align-items-start">
                            <div>
                                <h6 class="card-title">${result.name}</h6>
                                <p class="card-text text-muted small">${result.description || 'لا يوجد وصف'}</p>
                                <div class="d-flex gap-2">
                                    <span class="badge bg-primary">${result.author}</span>
                                    <span class="badge bg-secondary">${result.downloads || 0} تحميل</span>
                                    <span class="badge bg-success">${result.likes || 0} إعجاب</span>
                                    <span class="badge bg-info">${result.pipeline_tag || 'unknown'}</span>
                                </div>
                            </div>
                            <button class="btn btn-sm btn-outline-primary" onclick="window.app.modelsManager.addModelFromSearch('${result.id}', '${result.name}', '${result.description || ''}', '${result.pipeline_tag || 'text'}')">
                                <i class="fas fa-plus"></i> إضافة
                            </button>
                        </div>
                    </div>
                </div>
            `).join('');
        }

        searchResults.style.display = 'block';
    }

    async addModelFromSearch(modelId, name, description, pipelineTag) {
        try {
            // Determine category and modality from pipeline tag
            let category = 'text';
            let modality = 'text';

            if (pipelineTag.includes('image') || pipelineTag.includes('vision')) {
                category = 'vision';
                modality = 'vision';
            } else if (pipelineTag.includes('audio') || pipelineTag.includes('speech')) {
                category = 'audio';
                modality = 'audio';
            }

            const modelInfo = {
                name: name,
                model_id: modelId,
                category: category,
                type: 'teacher',
                description: description,
                modality: modality,
                architecture: 'transformer'
            };

            const success = await this.submitModel(modelInfo);
            if (success) {
                this.app.showSuccess(`تم إضافة النموذج: ${name}`);
                this.loadConfiguredModels();
            }

        } catch (error) {
            console.error('Error adding model from search:', error);
            this.app.showError('فشل في إضافة النموذج');
        }
    }

    async submitModel(modelInfo) {
        try {
            const response = await fetch('/api/models/add', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify(modelInfo)
            });

            const data = await response.json();
            return data.success;

        } catch (error) {
            console.error('Error submitting model:', error);
            this.app.showError('فشل في إضافة النموذج');
            return false;
        }
    }

    async toggleModelSelection(modelId, selected) {
        try {
            if (selected) {
                // Add to selected teachers
                if (!this.app.selectedTeachers.includes(modelId)) {
                    this.app.selectedTeachers.push(modelId);
                }
            } else {
                // Remove from selected teachers
                const index = this.app.selectedTeachers.indexOf(modelId);
                if (index > -1) {
                    this.app.selectedTeachers.splice(index, 1);
                }
            }

            // Update server
            const response = await fetch('/api/models/select', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({
                    teacher_models: this.app.selectedTeachers
                })
            });

            if (response.ok) {
                this.app.showSuccess(selected ? 'تم تحديد النموذج' : 'تم إلغاء تحديد النموذج');
                this.updateSelectedModelsDisplay();
            }

        } catch (error) {
            console.error('Error toggling model selection:', error);
            this.app.showError('فشل في تحديث اختيار النموذج');
        }
    }

    updateSelectedModelsDisplay() {
        // Update the selected models count and display
        const countElement = document.getElementById('model-count');
        if (countElement) {
            countElement.textContent = this.app.selectedTeachers.length;
        }

        // Update next step button
        const nextButton = document.getElementById('next-step-1');
        if (nextButton) {
            nextButton.disabled = this.app.selectedTeachers.length === 0;
        }

        // Update models grid display
        this.displaySelectedModels();
    }

    displaySelectedModels() {
        const modelsGrid = document.getElementById('models-grid');
        if (!modelsGrid) return;

        if (this.app.selectedTeachers.length === 0) {
            modelsGrid.innerHTML = '<p class="text-muted">لم يتم اختيار أي نماذج بعد</p>';
            return;
        }

        modelsGrid.innerHTML = this.app.selectedTeachers.map(modelId => {
            const model = this.app.configuredModels[modelId];
            if (!model) return '';

            return `
                <div class="model-card">
                    <div class="model-info">
                        <h6>${model.name}</h6>
                        <p class="text-muted small">${model.description || 'لا يوجد وصف'}</p>
                        <div class="model-badges">
                            <span class="badge bg-primary">${model.category}</span>
                            <span class="badge bg-secondary">${model.modality}</span>
                        </div>
                    </div>
                    <button class="btn btn-sm btn-outline-danger" onclick="window.app.modelsManager.toggleModelSelection('${modelId}', false)">
                        <i class="fas fa-times"></i>
                    </button>
                </div>
            `;
        }).join('');
    }

    async removeModel(modelId) {
        if (!confirm('هل أنت متأكد من حذف النموذج؟')) {
            return;
        }

        try {
            const response = await fetch(`/api/models/${encodeURIComponent(modelId)}`, {
                method: 'DELETE'
            });

            const data = await response.json();

            if (data.success) {
                this.app.showSuccess('تم حذف النموذج');
                this.loadConfiguredModels();
            } else {
                this.app.showError('فشل في حذف النموذج');
            }

        } catch (error) {
            console.error('Error removing model:', error);
            this.app.showError('خطأ في حذف النموذج');
        }
    }

    showModelInfo(modelId) {
        const model = this.app.configuredModels[modelId];
        if (model) {
            this.app.showInfo(`معلومات النموذج: ${model.name}\nالوصف: ${model.description}\nالفئة: ${model.category}\nالحجم: ${model.size}`);
        }
    }

    showAddCustomModelModal() {
        // Show modal for adding custom model
        this.app.showInfo('سيتم إضافة نافذة إضافة نموذج مخصص قريباً');
    }
}