document.addEventListener('DOMContentLoaded', () => {
    // DOM Elements
    const tabs = document.querySelectorAll('.tab-btn');
    const tabContents = document.querySelectorAll('.tab-content');
    const configForm = document.getElementById('config-form');
    
    // Toggles and fields
    const splitEnabled = document.getElementById('split-enabled');
    const splitRatio = document.getElementById('split-ratio');
    const splitClasses = document.getElementById('split-classes');
    
    // Dataset selectors
    const datasetSelect = document.getElementById('dataset-select');
    const datasetCustomWrapper = document.getElementById('dataset-custom-wrapper');
    
    // Model selectors
    const modelSelect = document.getElementById('model-select');
    const modelCustomWrapper = document.getElementById('model-custom-wrapper');
    
    const augmentationEnabled = document.getElementById('augmentation-enabled');
    const augmentationInputs = document.querySelectorAll('.augmentation-fields input, .augmentation-fields select');
    
    const mlflowEnabled = document.getElementById('mlflow-enabled');
    const mlflowInputs = document.querySelectorAll('.mlflow-fields input');
    const trackingUriInput = document.getElementById('tracking-uri');
    const mlflowHeaderLink = document.getElementById('mlflow-header-link');

    // --- Dynamic Model Selection ---
    const standardModels = {
        detect: ['yolo11n.pt', 'yolo11s.pt', 'yolo11m.pt', 'yolo11l.pt', 'yolo11x.pt'],
        segment: ['yolo11n-seg.pt', 'yolo11s-seg.pt', 'yolo11m-seg.pt', 'yolo11l-seg.pt', 'yolo11x-seg.pt'],
        classify: ['yolo11n-cls.pt', 'yolo11s-cls.pt', 'yolo11m-cls.pt', 'yolo11l-cls.pt', 'yolo11x-cls.pt'],
        pose: ['yolo11n-pose.pt', 'yolo11s-pose.pt', 'yolo11m-pose.pt', 'yolo11l-pose.pt', 'yolo11x-pose.pt'],
        obb: ['yolo11n-obb.pt', 'yolo11s-obb.pt', 'yolo11m-obb.pt', 'yolo11l-obb.pt', 'yolo11x-obb.pt']
    };
    let discoveredModels = [];

    function getDiscoveredModelValue(model) {
        return model.path || model.name;
    }

    function updateModelOptions() {
        const task = taskSelect.value;
        const stdModels = standardModels[task] || [];
        const currentSelectVal = modelSelect.value;
        
        modelSelect.innerHTML = '';
        
        // Group 1: Standard Models
        const stdGroup = document.createElement('optgroup');
        stdGroup.label = 'Стандартные модели';
        stdModels.forEach(model => {
            const opt = document.createElement('option');
            opt.value = model;
            opt.textContent = model;
            stdGroup.appendChild(opt);
        });
        modelSelect.appendChild(stdGroup);
        
        // Group 2: Discovered Models
        const localModels = discoveredModels.filter(m => getDiscoveredModelValue(m));
        if (localModels.length > 0) {
            const localGroup = document.createElement('optgroup');
            localGroup.label = 'Локальные/скачанные модели';
            localModels.forEach(m => {
                const opt = document.createElement('option');
                opt.value = getDiscoveredModelValue(m);
                opt.textContent = m.name;
                localGroup.appendChild(opt);
            });
            modelSelect.appendChild(localGroup);
        }
        
        // Custom Option
        const customOpt = document.createElement('option');
        customOpt.value = '__custom__';
        customOpt.textContent = 'Указать модель вручную...';
        modelSelect.appendChild(customOpt);
        
        // Match selection if valid
        const allAvailable = [...stdModels, ...localModels.map(getDiscoveredModelValue)];
        if (currentSelectVal === '__custom__' || allAvailable.includes(currentSelectVal)) {
            modelSelect.value = currentSelectVal;
        } else {
            modelSelect.value = stdModels[0] || '__custom__';
        }
        
        updateModelFieldsState();
    }

    function updateModelFieldsState() {
        const val = modelSelect.value;
        const modelInput = document.getElementById('model');
        if (val === '__custom__') {
            modelCustomWrapper.style.display = 'block';
        } else {
            modelCustomWrapper.style.display = 'none';
            modelInput.value = val;
        }
    }
    
    const taskSelect = document.getElementById('task');
    
    // Session Controls
    const sessionSelect = document.getElementById('session-select');
    const sessionNameInput = document.getElementById('session-name');
    const sessionSaveBtn = document.getElementById('session-save-btn');
    const sessionDeleteBtn = document.getElementById('session-delete-btn');
    
    // Control elements
    const startBtn = document.getElementById('start-btn');
    const stopBtn = document.getElementById('stop-btn');
    
    // Status elements
    const statusCard = document.getElementById('status-card');
    const statusTitle = document.getElementById('status-title');
    const statusText = document.getElementById('status-text');
    const statusTimer = document.getElementById('status-timer');
    const progressBarFill = document.getElementById('progress-bar-fill');
    const progressText = document.getElementById('progress-text');
    const progressEta = document.getElementById('progress-eta');
    
    // Logs
    const logContainer = document.getElementById('log-container');
    const autoscrollCheck = document.getElementById('autoscroll');
    const clearLogBtn = document.getElementById('clear-log-btn');
    
    // Chart
    const ctx = document.getElementById('metricsChart').getContext('2d');
    let metricsChart = null;
    
    // State variables
    let trainingTimer = null;
    let secondsElapsed = 0;
    let socket = null;
    let socketReconnectTimer = null;
    let isTrainingActive = false;
    let isTrainingStartPending = false;
    let trainingStatusRevision = 0;
    let currentTrainingStatus = 'idle';

    // --- View Tab Switching ---
    const viewBtns = document.querySelectorAll('.view-btn');
    const views = document.querySelectorAll('.workspace-view');

    viewBtns.forEach(btn => {
        btn.addEventListener('click', () => {
            viewBtns.forEach(b => b.classList.remove('active'));
            views.forEach(v => {
                v.classList.remove('active');
            });

            btn.classList.add('active');
            const viewId = btn.dataset.view;
            const view = document.getElementById(viewId);
            view.classList.add('active');
            
            localStorage.setItem('active_workspace_view', viewId);
        });
    });

    const savedView = localStorage.getItem('active_workspace_view');
    if (savedView) {
        const viewBtn = Array.from(viewBtns).find(b => b.dataset.view === savedView);
        if (viewBtn) viewBtn.click();
    }

    // --- Tab Switching ---
    tabs.forEach(tab => {
        tab.addEventListener('click', () => {
            tabs.forEach(t => t.classList.remove('active'));
            tabContents.forEach(c => c.classList.remove('active'));
            
            tab.classList.add('active');
            // Check if it's a sub-tab (has data-tab)
            if (tab.dataset.tab) {
                const contentId = `tab-${tab.dataset.tab}`;
                const content = document.getElementById(contentId);
                if (content) content.classList.add('active');
                localStorage.setItem('active_tab', tab.dataset.tab);
            }
        });
    });

    // Restore active tab on load
    const savedTab = localStorage.getItem('active_tab');
    if (savedTab) {
        const tabBtn = Array.from(tabs).find(t => t.dataset.tab === savedTab);
        if (tabBtn) {
            tabBtn.click();
        }
    }

    // --- Toggles & Constraints ---
    function updateSplitFields() {
        const isClassify = taskSelect.value === 'classify';
        if (isClassify && splitEnabled.checked) {
            splitEnabled.checked = false;
            showNotification('Для classify укажите готовый каталог с train/val по классам.', 'warning');
        }
        splitEnabled.disabled = isClassify;
        
        const disabled = !splitEnabled.checked || isClassify;
        splitRatio.disabled = disabled;
        splitClasses.disabled = disabled;
    }

    function updateAugmentationFields() {
        const disabled = !augmentationEnabled.checked;
        augmentationInputs.forEach(input => {
            input.disabled = disabled;
        });
    }

    function updateMlflowFields() {
        const disabled = !mlflowEnabled.checked;
        mlflowInputs.forEach(input => {
            input.disabled = disabled;
        });
        updateMlflowHeaderLink();
    }

    function updateMlflowHeaderLink() {
        const uri = trackingUriInput.value.trim();
        if (uri.startsWith('http://') || uri.startsWith('https://')) {
            let browserUri = uri;
            try {
                const parsedUri = new URL(uri);
                if (parsedUri.hostname === 'mlflow') {
                    parsedUri.hostname = window.location.hostname || 'localhost';
                    browserUri = parsedUri.href;
                }
            } catch (error) {
                console.error('Invalid MLflow tracking URI:', error);
            }
            mlflowHeaderLink.href = browserUri;
            mlflowHeaderLink.style.opacity = '1';
            mlflowHeaderLink.style.pointerEvents = 'auto';
        } else {
            mlflowHeaderLink.href = 'http://localhost:5000';
            mlflowHeaderLink.style.opacity = '0.5';
        }
    }

    trackingUriInput.addEventListener('input', updateMlflowHeaderLink);

    splitEnabled.addEventListener('change', updateSplitFields);
    taskSelect.addEventListener('change', () => {
        updateSplitFields();
        updateModelOptions();
    });
    augmentationEnabled.addEventListener('change', updateAugmentationFields);
    mlflowEnabled.addEventListener('change', updateMlflowFields);

    // --- Logger ---
    function addLogLine(message, level = 'info') {
        const line = document.createElement('div');
        line.className = `log-line log-level-${level.toLowerCase()}`;
        line.textContent = message;
        logContainer.appendChild(line);
        
        if (autoscrollCheck.checked) {
            logContainer.scrollTop = logContainer.scrollHeight;
        }
    }

    clearLogBtn.addEventListener('click', () => {
        logContainer.innerHTML = '';
    });

    // --- Chart.js Integration ---
    function initChart(datasets = []) {
        if (metricsChart) {
            metricsChart.destroy();
        }
        
        metricsChart = new Chart(ctx, {
            type: 'line',
            data: {
                labels: [],
                datasets: datasets
            },
            options: {
                responsive: true,
                maintainAspectRatio: false,
                scales: {
                    x: {
                        title: { display: true, text: 'Эпоха', color: '#a1a1aa' },
                        grid: { color: '#27272a' },
                        ticks: { color: '#a1a1aa' }
                    },
                    y: {
                        title: { display: true, text: 'Значение', color: '#a1a1aa' },
                        grid: { color: '#27272a' },
                        ticks: { color: '#a1a1aa' }
                    }
                },
                plugins: {
                    legend: {
                        labels: { color: '#f4f4f5', font: { family: 'Outfit' } }
                    }
                }
            }
        });
    }

    function updateChart(epoch, metrics) {
        if (!metricsChart) {
            initChart();
        }

        let labelIndex = metricsChart.data.labels.indexOf(epoch);
        if (labelIndex === -1) {
            metricsChart.data.labels.push(epoch);
            labelIndex = metricsChart.data.labels.length - 1;
            metricsChart.data.datasets.forEach(dataset => dataset.data.push(null));
        }

        const colors = ['#f97316', '#10b981', '#3b82f6', '#eab308', '#a855f7'];
        Object.entries(metrics).forEach(([key, value]) => {
            if (key === 'epoch') return;

            let dataset = metricsChart.data.datasets.find(item => item.label === key);
            if (!dataset) {
                const color = colors[metricsChart.data.datasets.length % colors.length];
                dataset = {
                    label: key,
                    data: Array(metricsChart.data.labels.length).fill(null),
                    borderColor: color,
                    backgroundColor: color + '22',
                    tension: 0.15,
                    fill: false
                };
                metricsChart.data.datasets.push(dataset);
            }

            dataset.data[labelIndex] = value;
        });

        metricsChart.update();
    }

    // --- Timer UI ---
    function startTimer() {
        stopTimer();
        secondsElapsed = 0;
        trainingTimer = setInterval(() => {
            secondsElapsed++;
            const h = String(Math.floor(secondsElapsed / 3600)).padStart(2, '0');
            const m = String(Math.floor((secondsElapsed % 3600) / 60)).padStart(2, '0');
            const s = String(secondsElapsed % 60).padStart(2, '0');
            statusTimer.textContent = `${h}:${m}:${s}`;
        }, 1000);
    }

    function stopTimer() {
        if (trainingTimer) {
            clearInterval(trainingTimer);
            trainingTimer = null;
        }
    }

    // --- WebSocket Sync ---
    function scheduleWebSocketReconnect() {
        if (socketReconnectTimer !== null) return;
        socketReconnectTimer = setTimeout(() => {
            socketReconnectTimer = null;
            connectWebSocket();
        }, 5000);
    }

    function parseWebSocketMessage(event, label, logLine) {
        try {
            const data = JSON.parse(event.data);
            if (!data || typeof data !== 'object' || Array.isArray(data)) {
                throw new Error('WebSocket payload must be an object');
            }
            return data;
        } catch (error) {
            console.error(`${label} message error:`, error);
            logLine('Получено некорректное сообщение от сервера.', 'warning');
            return null;
        }
    }

    function connectWebSocket() {
        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
        const wsUrl = `${protocol}//${window.location.host}/api/ws`;

        let nextSocket;
        try {
            nextSocket = new WebSocket(wsUrl);
        } catch (error) {
            console.error('WS connection error:', error);
            addLogLine('Не удалось подключиться к серверу. Повторная попытка через 5 секунд...', 'warning');
            scheduleWebSocketReconnect();
            return;
        }
        socket = nextSocket;
        
        nextSocket.onopen = () => {
            addLogLine('Соединение с сервером установлено.', 'info');
        };
        
        nextSocket.onclose = () => {
            if (socket !== nextSocket) return;
            socket = null;
            addLogLine('Соединение потеряно. Повторная попытка через 5 секунд...', 'warning');
            scheduleWebSocketReconnect();
        };
        
        nextSocket.onerror = (err) => {
            console.error('WS Error:', err);
        };
        
        nextSocket.onmessage = (event) => {
            const data = parseWebSocketMessage(event, 'Training WebSocket', addLogLine);
            if (!data) return;
            
            if (data.type === 'init') {
                updateUIStatus(data.status);
                
                // Load logs
                logContainer.innerHTML = '';
                (Array.isArray(data.logs) ? data.logs : []).forEach(([levelCode, msg]) => {
                    const level = levelCode.replace('__LOG_LEVEL_', '').replace('__', '').toLowerCase();
                    addLogLine(msg, level);
                });
                
                // Draw initial chart points
                initChart();
                if (Array.isArray(data.metrics) && data.metrics.length > 0) {
                    data.metrics.forEach(m => {
                        updateChart(m.epoch, m);
                    });
                }
                
                // Sync progress
                if (data.status === 'training' || data.status === 'stopping') {
                    updateProgress(data.epoch, data.total_epochs);
                }
            } else if (data.type === 'status') {
                updateUIStatus(data.status);
                if (data.output_dir) {
                    addLogLine(`Результаты сохранены: ${data.output_dir}`, 'success');
                }
            } else if (data.type === 'log') {
                const autoscrollCheck = document.getElementById('autoscroll');
                if (data.level === 'progress') {
                    let lastLine = logContainer.lastElementChild;
                    if (lastLine && lastLine.classList.contains('log-line-progress')) {
                        lastLine.textContent = data.message;
                    } else {
                        const line = document.createElement('div');
                        line.className = 'log-line log-line-progress log-level-info';
                        line.textContent = data.message;
                        logContainer.appendChild(line);
                    }
                } else {
                    let lastLine = logContainer.lastElementChild;
                    if (lastLine && lastLine.classList.contains('log-line-progress')) {
                        lastLine.classList.remove('log-line-progress');
                    }
                    addLogLine(data.message, data.level);
                }
                if (autoscrollCheck && autoscrollCheck.checked) {
                    logContainer.scrollTop = logContainer.scrollHeight;
                }
            } else if (data.type === 'progress') {
                updateProgress(data.epoch, data.total_epochs, data.message);
                if (data.metrics) {
                    updateChart(data.epoch, data.metrics);
                }
            }
        };
    }

    function updateUIStatus(status) {
        trainingStatusRevision++;
        currentTrainingStatus = status;
        statusCard.className = `status-${status}`;
        
        switch (status) {
            case 'idle':
                statusTitle.textContent = 'ГОТОВО К ЗАПУСКУ';
                statusText.textContent = 'Проверьте параметры и начните обучение.';
                isTrainingActive = false;
                startBtn.disabled = false;
                stopBtn.disabled = true;
                stopTimer();
                break;
            case 'preparing':
                statusTitle.textContent = 'ПОДГОТОВКА';
                statusText.textContent = 'Загрузка модели, разметки и настройка окружения...';
                isTrainingActive = true;
                startBtn.disabled = true;
                stopBtn.disabled = false;
                startTimer();
                initChart();
                break;
            case 'training':
                statusTitle.textContent = 'ОБУЧЕНИЕ';
                isTrainingActive = true;
                startBtn.disabled = true;
                stopBtn.disabled = false;
                if (!trainingTimer) startTimer();
                break;
            case 'stopping':
                statusTitle.textContent = 'ОСТАНОВКА';
                statusText.textContent = 'Остановка процессов обучения. Дождитесь закрытия...';
                isTrainingActive = true;
                startBtn.disabled = true;
                stopBtn.disabled = true;
                break;
            case 'finished': // Compatibility with sessions created by older versions.
            case 'succeeded':
                statusTitle.textContent = 'ГОТОВО';
                statusText.textContent = 'Обучение успешно завершено.';
                isTrainingActive = false;
                startBtn.disabled = false;
                stopBtn.disabled = true;
                stopTimer();
                break;
            case 'cancelled':
                statusTitle.textContent = 'ОСТАНОВЛЕНО';
                statusText.textContent = 'Обучение остановлено пользователем.';
                isTrainingActive = false;
                startBtn.disabled = false;
                stopBtn.disabled = true;
                stopTimer();
                break;
            case 'failed':
                statusTitle.textContent = 'ОШИБКА';
                statusText.textContent = 'Процесс завершился с ошибкой. Проверьте логи.';
                isTrainingActive = false;
                startBtn.disabled = false;
                stopBtn.disabled = true;
                stopTimer();
                break;
        }
    }

    function updateProgress(epoch, total, message = '') {
        const percent = total > 0 ? (epoch / total) * 100 : 0;
        progressBarFill.style.width = `${percent}%`;
        progressText.textContent = `Эпохи: ${epoch} / ${total}`;
        
        if (message) {
            statusText.textContent = message;
        }
    }

    // --- Read/Write Configurations ---
    function readNumber(id, fallback, integer = false) {
        const rawValue = document.getElementById(id).value;
        const value = integer
            ? Number.parseInt(rawValue, 10)
            : Number.parseFloat(rawValue);
        return Number.isNaN(value) ? fallback : value;
    }

    function readValidatedNumber(id, label, {integer = false, min, max} = {}) {
        const rawValue = document.getElementById(id).value.trim();
        const value = Number(rawValue);
        const outOfRange = (min !== undefined && value < min)
            || (max !== undefined && value > max);
        if (
            rawValue === ''
            || !Number.isFinite(value)
            || (integer && !Number.isInteger(value))
            || outOfRange
        ) {
            const range = max === undefined ? `не меньше ${min}` : `от ${min} до ${max}`;
            const integerHint = integer ? 'целым числом ' : '';
            showNotification(`${label} должен быть ${integerHint}${range}.`, 'warning');
            return null;
        }
        return value;
    }

    function getFormConfig() {
        return {
            dataset: document.getElementById('dataset').value.trim(),
            model: document.getElementById('model').value.trim(),
            task: taskSelect.value,
            epochs: readNumber('epochs', 100, true),
            image_size: readNumber('image-size', 640, true),
            batch_size: readNumber('batch-size', 16, true),
            device: document.getElementById('device').value.trim(),
            workers: readNumber('workers', 8, true),
            patience: readNumber('patience', 100, true),
            project: document.getElementById('project').value.trim() || 'runs/train',
            run_name: document.getElementById('run-name').value.trim(),
            split: {
                enabled: splitEnabled.checked,
                train_ratio: readNumber('split-ratio', 0.8),
                classes_path: splitClasses.value.trim()
            },
            augmentation: {
                enabled: augmentationEnabled.checked,
                hsv_h: readNumber('hsv-h', 0.015),
                hsv_s: readNumber('hsv-s', 0.7),
                hsv_v: readNumber('hsv-v', 0.4),
                degrees: readNumber('degrees', 0),
                translate: readNumber('translate', 0.1),
                scale: readNumber('scale', 0.5),
                shear: readNumber('shear', 0),
                perspective: readNumber('perspective', 0),
                close_mosaic: readNumber('close-mosaic', 10, true),
                flipud: readNumber('flipud', 0),
                fliplr: readNumber('fliplr', 0.5),
                bgr: readNumber('bgr', 0),
                mosaic: readNumber('mosaic', 1),
                mixup: readNumber('mixup', 0),
                cutmix: readNumber('cutmix', 0),
                copy_paste: readNumber('copy-paste', 0),
                erasing: readNumber('erasing', 0.4),
                copy_paste_mode: document.getElementById('copy-paste-mode').value,
                auto_augment: document.getElementById('auto-augment').value
            },
            mlflow: {
                enabled: mlflowEnabled.checked,
                tracking_uri: document.getElementById('tracking-uri').value.trim(),
                experiment_name: document.getElementById('experiment-name').value.trim(),
                run_name: document.getElementById('mlflow-run-name').value.trim()
            }
        };
    }

    function applyConfig(data) {
        taskSelect.value = data.task || 'detect';
        updateModelOptions();
        
        const modelVal = data.model || 'yolo11n.pt';
        const task = data.task || 'detect';
        const stdModels = standardModels[task] || [];
        const matchedLocalModel = discoveredModels.find(
            model => model.path === modelVal || model.name === modelVal
        );
        if (stdModels.includes(modelVal)) {
            modelSelect.value = modelVal;
            modelCustomWrapper.style.display = 'none';
            document.getElementById('model').value = modelVal;
        } else if (matchedLocalModel) {
            const discoveredValue = getDiscoveredModelValue(matchedLocalModel);
            modelSelect.value = discoveredValue;
            modelCustomWrapper.style.display = 'none';
            document.getElementById('model').value = discoveredValue;
        } else {
            modelSelect.value = '__custom__';
            modelCustomWrapper.style.display = 'block';
            document.getElementById('model').value = modelVal;
        }
        
        const matchedDataset = discoveredDatasets.find(d => d.path === data.dataset);
        if (matchedDataset) {
            datasetSelect.value = data.dataset;
            datasetCustomWrapper.style.display = 'none';
            document.getElementById('dataset').value = data.dataset;
        } else {
            datasetSelect.value = '__custom__';
            datasetCustomWrapper.style.display = 'block';
            document.getElementById('dataset').value = data.dataset || '';
        }
        
        // Split
        splitEnabled.checked = data.split?.enabled ?? false;
        splitRatio.value = data.split?.train_ratio ?? 0.8;
        splitClasses.value = data.split?.classes_path || '';
        
        // Training params
        document.getElementById('epochs').value = data.epochs ?? 100;
        document.getElementById('image-size').value = data.image_size ?? 640;
        document.getElementById('batch-size').value = data.batch_size ?? 16;
        document.getElementById('device').value = data.device || '';
        document.getElementById('workers').value = data.workers ?? 8;
        document.getElementById('patience').value = data.patience ?? 100;
        document.getElementById('project').value = data.project || 'runs/train';
        document.getElementById('run-name').value = data.run_name || '';
        
        // Augmentation
        augmentationEnabled.checked = data.augmentation?.enabled !== false;
        if (data.augmentation) {
            document.getElementById('hsv-h').value = data.augmentation.hsv_h ?? 0.015;
            document.getElementById('hsv-s').value = data.augmentation.hsv_s ?? 0.7;
            document.getElementById('hsv-v').value = data.augmentation.hsv_v ?? 0.4;
            document.getElementById('degrees').value = data.augmentation.degrees ?? 0.0;
            document.getElementById('translate').value = data.augmentation.translate ?? 0.1;
            document.getElementById('scale').value = data.augmentation.scale ?? 0.5;
            document.getElementById('shear').value = data.augmentation.shear ?? 0.0;
            document.getElementById('perspective').value = data.augmentation.perspective ?? 0.0;
            document.getElementById('close-mosaic').value = data.augmentation.close_mosaic ?? 10;
            document.getElementById('flipud').value = data.augmentation.flipud ?? 0.0;
            document.getElementById('fliplr').value = data.augmentation.fliplr ?? 0.5;
            document.getElementById('bgr').value = data.augmentation.bgr ?? 0.0;
            document.getElementById('mosaic').value = data.augmentation.mosaic ?? 1.0;
            document.getElementById('mixup').value = data.augmentation.mixup ?? 0.0;
            document.getElementById('cutmix').value = data.augmentation.cutmix ?? 0.0;
            document.getElementById('copy-paste').value = data.augmentation.copy_paste ?? 0.0;
            document.getElementById('erasing').value = data.augmentation.erasing ?? 0.4;
            document.getElementById('copy-paste-mode').value = data.augmentation.copy_paste_mode || 'flip';
            document.getElementById('auto-augment').value = data.augmentation.auto_augment || 'randaugment';
        }
        
        // MLflow
        mlflowEnabled.checked = data.mlflow?.enabled !== false;
        if (data.mlflow) {
            document.getElementById('tracking-uri').value = data.mlflow.tracking_uri || 'sqlite:///mlflow.db';
            document.getElementById('experiment-name').value = data.mlflow.experiment_name || 'yolo-webui';
            document.getElementById('mlflow-run-name').value = data.mlflow.run_name || '';
        }

        // Sync disables
        updateSplitFields();
        updateAugmentationFields();
        updateMlflowFields();
    }

    // --- Load Configuration (Last Run or Defaults) ---
    async function loadInitialConfig() {
        // 1. Try loading draft configuration from localStorage
        const draft = localStorage.getItem('draft_config');
        if (draft) {
            try {
                const data = JSON.parse(draft);
                applyConfig(data);
                addLogLine('Восстановлены последние измененные параметры.', 'info');
                return;
            } catch (e) {
                // Ignore and fall back
            }
        }

        // 2. First check if last_run exists
        try {
            const lastRes = await fetch('/api/sessions/last_run');
            if (lastRes.ok) {
                const data = await lastRes.json();
                applyConfig(data);
                addLogLine('Загружена конфигурация последнего запуска.', 'info');
                return;
            }
        } catch (e) {
            // Silence fail to fall back to defaults
        }

        // 3. Fall back to defaults
        try {
            const res = await fetch('/api/config/defaults');
            if (!res.ok) throw new Error('Failed to fetch defaults');
            const data = await res.json();
            applyConfig(data);
        } catch (err) {
            console.error('Error loading defaults:', err);
            showNotification('Ошибка загрузки настроек по умолчанию', 'error');
        }
    }

    // --- Sessions Management ---
    async function loadSessionsList() {
        try {
            const res = await fetch('/api/sessions');
            if (!res.ok) throw new Error();
            const names = await res.json();
            
            // Re-populate select
            const currentValue = sessionSelect.value;
            sessionSelect.innerHTML = '<option value="">По умолчанию (Последний запуск)</option>';
            names.forEach(name => {
                const opt = document.createElement('option');
                opt.value = name;
                opt.textContent = name;
                sessionSelect.appendChild(opt);
            });
            
            // Restore selection if still exists
            const savedProfile = localStorage.getItem('selected_profile') || "";
            const finalValue = currentValue || savedProfile;
            if (names.includes(finalValue)) {
                sessionSelect.value = finalValue;
                sessionDeleteBtn.disabled = false;
            } else {
                sessionSelect.value = "";
                sessionDeleteBtn.disabled = true;
            }
        } catch (e) {
            console.error("Failed to load sessions list:", e);
        }
    }

    // --- Datasets Auto-Discovery ---
    let discoveredDatasets = [];

    async function loadDatasetsList() {
        try {
            const res = await fetch('/api/datasets');
            if (!res.ok) throw new Error();
            discoveredDatasets = await res.json();
            
            // Re-populate select
            datasetSelect.innerHTML = '';
            discoveredDatasets.forEach(item => {
                const opt = document.createElement('option');
                opt.value = item.path;
                opt.textContent = `${item.name} (${item.type === 'directory' ? 'Папка' : 'Конфиг'})`;
                datasetSelect.appendChild(opt);
            });
            
            // Add custom path option
            const customOpt = document.createElement('option');
            customOpt.value = '__custom__';
            customOpt.textContent = 'Указать путь вручную...';
            datasetSelect.appendChild(customOpt);
            
            updateDatasetFieldsState();
        } catch (e) {
            console.error("Failed to load datasets list:", e);
            datasetSelect.innerHTML = '<option value="__custom__">Указать путь вручную...</option>';
            updateDatasetFieldsState();
        }
    }

    function updateDatasetFieldsState() {
        const val = datasetSelect.value;
        const datasetInput = document.getElementById('dataset');
        
        if (val === '__custom__') {
            datasetCustomWrapper.style.display = 'block';
        } else {
            datasetCustomWrapper.style.display = 'none';
            datasetInput.value = val;
        }
    }

    datasetSelect.addEventListener('change', updateDatasetFieldsState);

    async function loadModelsList() {
        try {
            const res = await fetch('/api/models');
            if (!res.ok) throw new Error();
            discoveredModels = await res.json();
            updateExportModelOptions();
        } catch (e) {
            console.error("Failed to load models list:", e);
        }
    }

    modelSelect.addEventListener('change', updateModelFieldsState);

    // --- Export Model Selection ---
    const exportModelSelect = document.getElementById('export-model-select');
    const exportModelCustomWrapper = document.getElementById('export-model-custom-wrapper');
    const exportModelInput = document.getElementById('export-model');

    function updateExportModelOptions() {
        exportModelSelect.innerHTML = '';
        
        if (discoveredModels.length > 0) {
            const localGroup = document.createElement('optgroup');
            localGroup.label = 'Локальные/обученные модели';
            discoveredModels.forEach(m => {
                const opt = document.createElement('option');
                opt.value = m.path;
                opt.textContent = m.source === 'runs' ? `[RUNS] ${m.name}` : m.name;
                localGroup.appendChild(opt);
            });
            exportModelSelect.appendChild(localGroup);
        }
        
        const customOpt = document.createElement('option');
        customOpt.value = '__custom__';
        customOpt.textContent = 'Указать путь вручную...';
        exportModelSelect.appendChild(customOpt);
        
        updateExportModelFieldsState();
    }

    function updateExportModelFieldsState() {
        const val = exportModelSelect.value;
        if (val === '__custom__') {
            exportModelCustomWrapper.style.display = 'block';
        } else {
            exportModelCustomWrapper.style.display = 'none';
            exportModelInput.value = val;
        }
    }

    exportModelSelect.addEventListener('change', updateExportModelFieldsState);

    sessionSelect.addEventListener('change', async () => {
        const name = sessionSelect.value;
        localStorage.setItem('selected_profile', name);
        if (!name) {
            sessionDeleteBtn.disabled = true;
            localStorage.removeItem('draft_config'); // Reset draft
            await loadInitialConfig();
            return;
        }
        
        sessionDeleteBtn.disabled = false;
        try {
            const res = await fetch(`/api/sessions/${name}`);
            if (!res.ok) throw new Error();
            const data = await res.json();
            applyConfig(data);
            localStorage.setItem('draft_config', JSON.stringify(data));
            showNotification(`Профиль "${name}" успешно загружен.`, 'success');
        } catch (e) {
            showNotification('Не удалось загрузить выбранный профиль.', 'error');
        }
    });

    sessionSaveBtn.addEventListener('click', async () => {
        const name = sessionNameInput.value.trim();
        if (!/^[a-zA-Z0-9_-]+$/.test(name)) {
            showNotification('Введите корректное имя профиля (латиница, цифры, дефисы и подчёркивания).', 'warning');
            return;
        }
        if (name === "last_run") {
            showNotification('Имя "last_run" зарезервировано бэкендом.', 'warning');
            return;
        }

        const config = getFormConfig();
        try {
            const res = await fetch(`/api/sessions/${name}`, {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(config)
            });
            if (!res.ok) throw new Error();
            
            showNotification(`Профиль "${name}" сохранен.`, 'success');
            sessionNameInput.value = "";
            
            localStorage.setItem('selected_profile', name);
            localStorage.setItem('draft_config', JSON.stringify(config));
            await loadSessionsList();
            sessionSelect.value = name;
            sessionDeleteBtn.disabled = false;
        } catch (e) {
            showNotification('Не удалось сохранить профиль.', 'error');
        }
    });

    sessionDeleteBtn.addEventListener('click', async () => {
        const name = sessionSelect.value;
        if (!name) return;
        
        if (!confirm(`Вы действительно хотите удалить профиль "${name}"?`)) return;

        try {
            const res = await fetch(`/api/sessions/${name}`, { method: 'DELETE' });
            if (!res.ok) throw new Error();
            
            showNotification(`Профиль "${name}" удален.`, 'success');
            sessionSelect.value = "";
            sessionDeleteBtn.disabled = true;
            localStorage.removeItem('selected_profile');
            localStorage.removeItem('draft_config');
            await loadSessionsList();
            await loadInitialConfig();
        } catch (e) {
            console.error('Delete profile error:', e);
            showNotification('Не удалось удалить профиль.', 'error');
        }
    });

    // --- Form submit ---
    async function startTraining() {
        if (isTrainingActive || isTrainingStartPending) return;
        isTrainingStartPending = true;
        startBtn.disabled = true;
        const config = getFormConfig();
        const statusRevisionAtStart = trainingStatusRevision;
        let started = false;

        try {
            const res = await fetch('/api/train/start', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(config)
            });
            
            const data = await res.json();
            if (!res.ok) {
                throw new Error(data.detail || 'Failed to start training');
            }
            started = true;
            const terminalStatusReceived = trainingStatusRevision !== statusRevisionAtStart
                && ['finished', 'succeeded', 'cancelled', 'failed'].includes(currentTrainingStatus);
            if (!isTrainingActive && !terminalStatusReceived) {
                updateUIStatus('preparing');
            }
            showNotification('Обучение успешно запущено!', 'success');
        } catch (err) {
            console.error('Start error:', err);
            showNotification(err.message, 'error');
        } finally {
            isTrainingStartPending = false;
            if (!started && !isTrainingActive) {
                startBtn.disabled = false;
            }
        }
    }

    async function stopTraining() {
        if (!isTrainingActive) return;
        try {
            const res = await fetch('/api/train/stop', { method: 'POST' });
            if (!res.ok) {
                const data = await res.json();
                throw new Error(data.detail || 'Failed to stop training');
            }
            showNotification('Запрос на остановку отправлен.', 'info');
        } catch (err) {
            console.error('Stop error:', err);
            showNotification(err.message, 'error');
        }
    }

    if (configForm) {
        configForm.addEventListener('input', () => {
            const config = getFormConfig();
            localStorage.setItem('draft_config', JSON.stringify(config));
        });
        configForm.addEventListener('change', () => {
            const config = getFormConfig();
            localStorage.setItem('draft_config', JSON.stringify(config));
        });
    }

    startBtn.addEventListener('click', startTraining);
    stopBtn.addEventListener('click', stopTraining);

    // --- Helper Notification System ---
    function showNotification(message, type = 'info') {
        const toast = document.createElement('div');
        toast.style.position = 'fixed';
        toast.style.bottom = '20px';
        toast.style.right = '20px';
        toast.style.padding = '12px 20px';
        toast.style.borderRadius = '8px';
        toast.style.fontFamily = 'Outfit';
        toast.style.fontSize = '0.9rem';
        toast.style.fontWeight = '500';
        toast.style.zIndex = '9999';
        toast.style.boxShadow = '0 10px 25px rgba(0,0,0,0.5)';
        toast.style.animation = 'slideIn 0.3s cubic-bezier(0.4, 0, 0.2, 1)';
        toast.style.maxWidth = '350px';
        
        if (type === 'success') {
            toast.style.backgroundColor = 'var(--success)';
            toast.style.color = '#000';
        } else if (type === 'error') {
            toast.style.backgroundColor = 'var(--error)';
            toast.style.color = '#fff';
        } else if (type === 'warning') {
            toast.style.backgroundColor = 'var(--warning)';
            toast.style.color = '#000';
        } else {
            toast.style.backgroundColor = 'var(--accent)';
            toast.style.color = '#fff';
        }
        
        toast.textContent = message;
        document.body.appendChild(toast);
        
        setTimeout(() => {
            toast.style.animation = 'fadeOut 0.5s ease forwards';
            setTimeout(() => toast.remove(), 500);
        }, 4000);
    }

    // Add keyframes dynamically if not in stylesheet
    const styleSheet = document.createElement("style");
    styleSheet.innerText = `
        @keyframes slideIn {
            from { transform: translateY(100%) scale(0.9); opacity: 0; }
            to { transform: translateY(0) scale(1); opacity: 1; }
        }
        @keyframes fadeOut {
            from { opacity: 1; }
            to { opacity: 0; }
        }
    `;
    document.head.appendChild(styleSheet);

    // --- Export Logic ---
    const exportStartBtn = document.getElementById('export-start-btn');
    const exportStopBtn = document.getElementById('export-stop-btn');
    const exportStatusCard = document.getElementById('export-status-card');
    const exportStatusTitle = document.getElementById('export-status-title');
    const exportStatusText = document.getElementById('export-status-text');
    const exportLogContainer = document.getElementById('export-log-container');
    const exportClearLogBtn = document.getElementById('export-clear-log-btn');
    const exportAutoscrollCheck = document.getElementById('export-autoscroll');
    
    let isExportActive = false;
    let isExportStartPending = false;
    let exportSocket = null;
    let exportSocketReconnectTimer = null;
    let exportStatusRevision = 0;
    let currentExportStatus = 'idle';

    function addExportLogLine(message, level = 'info') {
        const line = document.createElement('div');
        line.className = `log-line log-level-${level.toLowerCase()}`;
        line.textContent = message;
        exportLogContainer.appendChild(line);
        
        if (exportAutoscrollCheck.checked) {
            exportLogContainer.scrollTop = exportLogContainer.scrollHeight;
        }
    }

    exportClearLogBtn.addEventListener('click', () => {
        exportLogContainer.innerHTML = '';
    });

    function scheduleExportWebSocketReconnect() {
        if (exportSocketReconnectTimer !== null) return;
        exportSocketReconnectTimer = setTimeout(() => {
            exportSocketReconnectTimer = null;
            connectExportWebSocket();
        }, 5000);
    }

    function connectExportWebSocket() {
        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
        const wsUrl = `${protocol}//${window.location.host}/api/export/ws`;

        let nextSocket;
        try {
            nextSocket = new WebSocket(wsUrl);
        } catch (error) {
            console.error('Export WS connection error:', error);
            addExportLogLine('Не удалось подключиться к серверу. Повторная попытка через 5 секунд...', 'warning');
            scheduleExportWebSocketReconnect();
            return;
        }
        exportSocket = nextSocket;
        
        nextSocket.onopen = () => {
            addExportLogLine('Соединение с сервером установлено.', 'info');
        };
        
        nextSocket.onclose = () => {
            if (exportSocket !== nextSocket) return;
            exportSocket = null;
            addExportLogLine('Соединение потеряно. Повторная попытка через 5 секунд...', 'warning');
            scheduleExportWebSocketReconnect();
        };
        
        nextSocket.onerror = (err) => {
            console.error('Export WS Error:', err);
        };
        
        nextSocket.onmessage = (event) => {
            const data = parseWebSocketMessage(event, 'Export WebSocket', addExportLogLine);
            if (!data) return;
            
            if (data.type === 'init') {
                updateExportUIStatus(data.status);
                exportLogContainer.innerHTML = '';
                (Array.isArray(data.logs) ? data.logs : []).forEach(([levelCode, msg]) => {
                    const level = levelCode.replace('__LOG_LEVEL_', '').replace('__', '').toLowerCase();
                    addExportLogLine(msg, level);
                });
            } else if (data.type === 'status') {
                updateExportUIStatus(data.status);
                if (data.output_dir) {
                    addExportLogLine(`Результаты сохранены: ${data.output_dir}`, 'success');
                }
            } else if (data.type === 'log') {
                addExportLogLine(data.message, data.level);
            }
        };
    }

    function updateExportUIStatus(status) {
        exportStatusRevision++;
        currentExportStatus = status;
        exportStatusCard.className = `status-${status}`;
        
        switch (status) {
            case 'idle':
                exportStatusTitle.textContent = 'ГОТОВО К ЭКСПОРТУ';
                exportStatusText.textContent = 'Выберите модель и формат для начала.';
                isExportActive = false;
                exportStartBtn.disabled = false;
                exportStopBtn.disabled = true;
                break;
            case 'preparing':
            case 'exporting':
                exportStatusTitle.textContent = 'ЭКСПОРТ';
                exportStatusText.textContent = 'Выполняется экспорт модели...';
                isExportActive = true;
                exportStartBtn.disabled = true;
                exportStopBtn.disabled = false;
                break;
            case 'stopping':
                exportStatusTitle.textContent = 'ОСТАНОВКА';
                exportStatusText.textContent = 'Остановка процесса экспорта...';
                isExportActive = true;
                exportStartBtn.disabled = true;
                exportStopBtn.disabled = true;
                break;
            case 'succeeded':
                exportStatusTitle.textContent = 'ГОТОВО';
                exportStatusText.textContent = 'Экспорт успешно завершен.';
                isExportActive = false;
                exportStartBtn.disabled = false;
                exportStopBtn.disabled = true;
                break;
            case 'cancelled':
                exportStatusTitle.textContent = 'ОСТАНОВЛЕНО';
                exportStatusText.textContent = 'Экспорт остановлен пользователем.';
                isExportActive = false;
                exportStartBtn.disabled = false;
                exportStopBtn.disabled = true;
                break;
            case 'failed':
                exportStatusTitle.textContent = 'ОШИБКА';
                exportStatusText.textContent = 'Экспорт завершился с ошибкой.';
                isExportActive = false;
                exportStartBtn.disabled = false;
                exportStopBtn.disabled = true;
                break;
        }
    }

    async function startExport() {
        if (isExportActive || isExportStartPending) return;
        
        const modelVal = document.getElementById('export-model').value.trim();
        if (!modelVal) {
            showNotification('Пожалуйста, выберите или укажите модель для экспорта.', 'warning');
            return;
        }

        const imgsz = readValidatedNumber(
            'export-imgsz',
            'Размер изображения',
            {integer: true, min: 32, max: 8192}
        );
        if (imgsz === null) return;
        const batch = readValidatedNumber(
            'export-batch',
            'Размер батча',
            {integer: true, min: 1, max: 1024}
        );
        if (batch === null) return;
        const workspace = readValidatedNumber(
            'export-workspace',
            'Workspace',
            {min: 1, max: 64}
        );
        if (workspace === null) return;

        const half = document.getElementById('export-half').checked;
        const int8 = document.getElementById('export-int8').checked;
        if (half && int8) {
            showNotification('FP16 и INT8 нельзя включать одновременно.', 'warning');
            return;
        }

        isExportStartPending = true;
        exportStartBtn.disabled = true;
        const config = {
            model: modelVal,
            format: document.getElementById('export-format').value,
            imgsz,
            half,
            int8,
            dynamic: document.getElementById('export-dynamic').checked,
            simplify: document.getElementById('export-simplify').checked,
            batch,
            workspace
        };
        const statusRevisionAtStart = exportStatusRevision;
        let started = false;

        try {
            const res = await fetch('/api/export/start', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(config)
            });
            
            const data = await res.json();
            if (!res.ok) {
                throw new Error(data.detail || 'Failed to start export');
            }
            started = true;
            const terminalStatusReceived = exportStatusRevision !== statusRevisionAtStart
                && ['succeeded', 'cancelled', 'failed'].includes(currentExportStatus);
            if (!isExportActive && !terminalStatusReceived) {
                updateExportUIStatus('preparing');
            }
            showNotification('Экспорт успешно запущен!', 'success');
        } catch (err) {
            console.error('Export start error:', err);
            showNotification(err.message, 'error');
        } finally {
            isExportStartPending = false;
            if (!started && !isExportActive) {
                exportStartBtn.disabled = false;
            }
        }
    }

    async function stopExport() {
        if (!isExportActive) return;
        try {
            const res = await fetch('/api/export/stop', { method: 'POST' });
            if (!res.ok) {
                const data = await res.json();
                throw new Error(data.detail || 'Failed to stop export');
            }
            showNotification('Запрос на остановку отправлен.', 'info');
        } catch (err) {
            console.error('Export stop error:', err);
            showNotification(err.message, 'error');
        }
    }

    exportStartBtn.addEventListener('click', startExport);
    exportStopBtn.addEventListener('click', stopExport);

    // Initial load sequence
    loadDatasetsList().then(() => {
        return loadModelsList();
    }).then(() => {
        return loadInitialConfig();
    }).then(() => {
        loadSessionsList();
        connectWebSocket();
        connectExportWebSocket();
        initChart();
    });
});