const { ipcRenderer } = require('electron');
const io = require('socket.io-client');

const socket = io('http://127.0.0.1:5001', {
    transports: ['websocket', 'polling'],
    upgrade: true,
    reconnection: true
});

const gridCells = [];
let latentHistory = [];
const entropyBars = [];

function logMsg(msg) {
    const logWindow = document.getElementById('log-window');
    if (!logWindow) return;
    const div = document.createElement('div');
    const time = new Date().toTimeString().split(' ')[0];
    div.innerText = `[${time}] ${msg}`;
    logWindow.appendChild(div);
    logWindow.scrollTop = logWindow.scrollHeight;
}

function updateFlux(z) {
    if (!gridCells.length) return;
    for (let i = 0; i < 128; i++) {
        const val = z[i] || 0;
        const b = Math.floor(Math.min(255, Math.max(0, (val + 2) * 42)));
        gridCells[i].style.background = `rgb(${b},${b},${b})`;
    }
}

function updateWaveform(z) {
    const canvas = document.getElementById('waveform-canvas');
    if (!canvas) return;
    const wctx = canvas.getContext('2d');
    if (canvas.width !== canvas.clientWidth) {
        canvas.width = canvas.clientWidth;
        canvas.height = canvas.clientHeight;
    }
    latentHistory.push(z[0]);
    if (latentHistory.length > 100) latentHistory.shift();
    wctx.clearRect(0, 0, canvas.width, canvas.height);
    wctx.beginPath();
    wctx.strokeStyle = '#E0E0E0';
    wctx.lineWidth = 1;
    for (let i = 0; i < latentHistory.length; i++) {
        const x = (i / 100) * canvas.width;
        const y = (canvas.height / 2) + (latentHistory[i] * 30);
        if (i === 0) wctx.moveTo(x, y);
        else wctx.lineTo(x, y);
    }
    wctx.stroke();
}

function updateEntropy(z) {
    if (!entropyBars.length) return;
    for (let i = 0; i < 32; i++) {
        const val = Math.abs(z[i*4] || 0);
        const height = Math.min(100, val * 40 + 5);
        entropyBars[i].style.height = `${height}%`;
        entropyBars[i].style.background = height > 70 ? '#E0E0E0' : '#333';
    }
}

function refreshModelList() {
    const modelSelector = document.getElementById('model-selector');
    if (!modelSelector) return;
    
    socket.emit('list_models', {}, (data) => {
        if (data && data.models) {
            modelSelector.innerHTML = '';
            data.models.forEach(model => {
                const opt = document.createElement('option');
                opt.value = model;
                opt.innerText = model;
                modelSelector.appendChild(opt);
            });
            logMsg(`REGISTRY: ${data.models.length} MODELS_FOUND`);
        }
    });
}

document.addEventListener('DOMContentLoaded', () => {
    const fluxGrid = document.getElementById('flux-grid');
    if (fluxGrid) {
        for (let i = 0; i < 128; i++) {
            const cell = document.createElement('div');
            cell.style.width = '100%';
            cell.style.height = '8px';
            cell.style.background = '#111';
            fluxGrid.appendChild(cell);
            gridCells.push(cell);
        }
    }

    const entropyGrid = document.getElementById('entropy-grid');
    if (entropyGrid) {
        for (let i = 0; i < 32; i++) {
            const bar = document.createElement('div');
            bar.style.flexGrow = '1';
            bar.style.height = '10%';
            bar.style.background = '#333';
            bar.style.transition = 'height 0.2s';
            entropyGrid.appendChild(bar);
            entropyBars.push(bar);
        }
    }

    const epIn = document.getElementById('epochs-input');
    const epOut = document.getElementById('epochs-display');
    if (epIn && epOut) epIn.oninput = function() { epOut.innerText = this.value; };

    const baIn = document.getElementById('batch-input');
    const baOut = document.getElementById('batch-display');
    if (baIn && baOut) baIn.oninput = function() { baOut.innerText = this.value; };

    const btnGen = document.getElementById('btn-gen');
    if (btnGen) {
        btnGen.onclick = () => {
            const count = parseInt(document.getElementById('gen-count').value) || 1000;
            logMsg(`HANDSHAKE_SENT: GEN_${count}`);
            socket.emit('start_generation', { count: count });
            btnGen.disabled = true;
            btnGen.innerText = "GENERATING...";
        };
    }

    const btnSave = document.getElementById('btn-save-model');
    const btnLoad = document.getElementById('btn-load-model');
    const modelSelector = document.getElementById('model-selector');
    const btnRefreshModels = document.getElementById('btn-refresh-models');

    if (modelSelector) {
        modelSelector.onchange = () => {
            const name = modelSelector.value;
            logMsg(`INIT_SWAP: ${name}`);
            socket.emit('select_model', { name: name });
        };
    }

    if (btnRefreshModels) {
        btnRefreshModels.onclick = refreshModelList;
    }

    if (btnSave) {
        btnSave.onclick = () => {
            const name = document.getElementById('model-save-name').value || 'archive';
            logMsg(`INIT_ARCHIVE_SEQUENCE: ${name}`);
            socket.emit('save_model', { name: name });
        };
    }

    if (btnLoad) {
        btnLoad.onclick = () => {
            const name = modelSelector.value;
            logMsg(`INIT_RELOAD_SEQUENCE: ${name}`);
            socket.emit('load_model', { name: name });
        };
    }

    const btnTrain = document.getElementById('btn-train');
    if (btnTrain) {
        btnTrain.onclick = () => {
            const epochs = parseInt(document.getElementById('epochs-input').value);
            const batch = parseInt(document.getElementById('batch-input').value);
            logMsg(`HANDSHAKE_SENT: TRAIN_E${epochs}`);
            socket.emit('start_training', { epochs: epochs, batch_size: batch });
            btnTrain.disabled = true;
            btnTrain.innerText = "TRAINING...";
            const canvas = document.getElementById('dream-canvas');
            if (canvas) canvas.classList.add('blur');
            const overlay = document.getElementById('training-overlay');
            if (overlay) overlay.style.display = 'flex';
        };
    }

    const btnPause = document.getElementById('btn-pause');
    const speedInput = document.getElementById('speed-input');
    const speedDisplay = document.getElementById('speed-display');
    let paused = false;

    if (btnPause) {
        btnPause.onclick = () => {
            paused = !paused;
            socket.emit('set_pause', { paused: paused });
            btnPause.innerText = paused ? "DREAM_STATE: PAUSED" : "DREAM_STATE: ACTIVE";
            btnPause.classList.toggle('active', paused);
            logMsg(`MOTION: ${paused ? 'PAUSED' : 'RESUMED'}`);
        };
    }

    if (speedInput && speedDisplay) {
        speedInput.oninput = function() {
            const val = parseFloat(this.value);
            speedDisplay.innerText = val.toFixed(3);
            socket.emit('set_speed', { speed: val });
        };
    }

    const winClose = document.getElementById('win-close');
    const winMin = document.getElementById('win-min');
    const winMax = document.getElementById('win-max');
    if (winClose) winClose.onclick = () => window.close();
    if (winMin) winMin.onclick = () => ipcRenderer.send('window-minimize');
    if (winMax) winMax.onclick = () => ipcRenderer.send('window-maximize');
});

socket.on('connect', () => {
    console.log("SOCKET_CONNECTED");
    logMsg("CORE_PIPE_SYNC: OK");
});

socket.on('model_status', (data) => {
    const modelStatus = document.getElementById('model-status');
    if (modelStatus) modelStatus.innerText = data.msg;
    logMsg(`MODEL_ARCHIVE: ${data.msg}`);
    if (data.msg.includes('SAVED')) refreshModelList();
});

let lastFrameTime = Date.now();
let frameCount = 0;

socket.on('frame', (data) => {
    const now = Date.now();
    const delta = (now - lastFrameTime) / 1000;
    lastFrameTime = now;
    frameCount++;
    if (frameCount % 60 === 0) logMsg(`PIPE_STREAM_SYNC [FPS: ${(1/delta).toFixed(1)}]`);

    const canvas = document.getElementById('dream-canvas');
    if (canvas) {
        const ctx = canvas.getContext('2d');
        const img = new Image();
        img.onload = () => {
            ctx.clearRect(0, 0, canvas.width, canvas.height);
            ctx.drawImage(img, 0, 0);
        };
        img.src = 'data:image/jpeg;base64,' + data.image;
    }
    
    updateFlux(data.z);
    updateWaveform(data.z);
    updateEntropy(data.z);

    const valSpeed = document.getElementById('val-speed');
    const latencyBar = document.getElementById('latency-bar');
    if (valSpeed) valSpeed.innerText = `${delta.toFixed(3)}s`;
    if (latencyBar) {
        const width = Math.min(100, delta * 1000);
        latencyBar.style.width = `${width}%`;
    }
});

socket.on('gen_status', (data) => {
    logMsg(`ENGINE: ${data.msg}`);
    const gs = document.getElementById('gen-status');
    if (gs) gs.innerText = data.msg;
});

socket.on('gen_complete', (data) => {
    logMsg(`SUCCESS: DATA_SYNTH_FINISH`);
    const btnGen = document.getElementById('btn-gen');
    if (btnGen) {
        btnGen.disabled = false;
        btnGen.innerText = "GENERATE_SAMPLES";
    }
});

socket.on('train_progress', (data) => {
    const ts = document.getElementById('train-status');
    if (ts) ts.innerText = `E:${data.epoch} | L:${data.loss.toFixed(4)}`;
    
    const pb = document.getElementById('progress-bar');
    if (pb) pb.style.width = `${data.progress}%`;

    const op = document.getElementById('overlay-percent');
    if (op) op.innerText = `${Math.floor(data.progress)}%`;

    const ol = document.getElementById('overlay-loss');
    if (ol) ol.innerText = `LOSS: ${data.loss.toFixed(6)}`;

    if (data.imagination) {
        const canvas = document.getElementById('dream-canvas');
        if (canvas) {
            const ctx = canvas.getContext('2d');
            const img = new Image();
            img.onload = () => {
                ctx.clearRect(0, 0, canvas.width, canvas.height);
                ctx.drawImage(img, 0, 0);
            };
            img.src = 'data:image/jpeg;base64,' + data.imagination;
        }
    }
});

socket.on('train_complete', (data) => {
    logMsg("TRAIN_COMPLETE: WEIGHTS_SYNCED");
    const btnTrain = document.getElementById('btn-train');
    if (btnTrain) {
        btnTrain.disabled = false;
        btnTrain.innerText = "START_TRAIN";
    }
    const canvas = document.getElementById('dream-canvas');
    if (canvas) canvas.classList.remove('blur');
    const overlay = document.getElementById('training-overlay');
    if (overlay) overlay.style.display = 'none';
    refreshModelList();
});

async function runBoot() {
    const bootScreen = document.getElementById('boot-screen');
    const mainUI = document.getElementById('main-ui');
    if (!bootScreen || !mainUI) return;

    const lines = [
        "DREAMER_ENGINE_V1.0 BIOS",
        "CONNECTING_TO_CORE_SVR...",
        "WAITING_FOR_HANDSHAKE..."
    ];
    for (let l of lines) {
        const d = document.createElement('div');
        d.innerText = l;
        bootScreen.appendChild(d);
        await new Promise(r => setTimeout(r, 200));
    }
    
    const forceStart = setTimeout(() => {
        if (bootScreen.style.display !== 'none') {
            console.log("FORCE_STARTING_UI");
            bootScreen.style.display = 'none';
            mainUI.style.display = 'flex';
            logMsg("SYSTEM_READY: EMERGENCY_BOOT");
            refreshModelList();
        }
    }, 5000);

    const wait = setInterval(() => {
        if (socket.connected) {
            clearInterval(wait);
            clearTimeout(forceStart);
            bootScreen.style.display = 'none';
            mainUI.style.display = 'flex';
            logMsg("SYSTEM_READY: PORT_5001");
            refreshModelList();
        }
    }, 500);
}

setInterval(() => {
    const clock = document.getElementById('sys-clock');
    if (clock) clock.innerText = new Date().toTimeString().split(' ')[0];
}, 1000);

runBoot();
