gpu-program-swapper/static/app.js

let evtSource = null;
let currentInstalledModels = [];

// Initialize SSE Stream
function initSSE() {
  if (evtSource) {
    evtSource.close();
  }
  
  evtSource = new EventSource('/api/stream');
  
  evtSource.onopen = () => {
    document.getElementById('sse-badge').textContent = 'SSE LIVE';
    document.getElementById('sse-badge').className = 'text-emerald-400 font-bold';
  };
  
  evtSource.onmessage = (event) => {
    try {
      const data = JSON.parse(event.data);
      updateDashboard(data);
    } catch (err) {
      console.error('Error parsing SSE event:', err);
    }
  };
  
  evtSource.onerror = (err) => {
    console.warn('SSE disconnected, retrying...', err);
    document.getElementById('sse-badge').textContent = 'RECONNECTING...';
    document.getElementById('sse-badge').className = 'text-rose-400 font-bold';
  };
}

function updateDashboard(data) {
  if (!data) return;

  // 1. GPU VRAM Stats
  const gpu = data.gpu || {};
  if (gpu.available) {
    document.getElementById('gpu-chip-name').textContent = gpu.device_name || 'NVIDIA GPU';
    document.getElementById('vram-total-used').textContent = gpu.vram_used_gb || '0.0';
    document.getElementById('vram-used-pct').textContent = `${gpu.vram_used_pct || 0}% USED`;
    
    const bd = gpu.breakdown || {};
    const totalBytes = gpu.vram_total_bytes || (16 * 1024**3);
    
    const ollamaPct = ((bd.ollama_gb || 0) / (gpu.vram_total_gb || 16)) * 100;
    const comfyPct = ((bd.comfyui_gb || 0) / (gpu.vram_total_gb || 16)) * 100;
    const systemPct = ((bd.system_gb || 0) / (gpu.vram_total_gb || 16)) * 100;
    const freePct = Math.max(0, 100 - (ollamaPct + comfyPct + systemPct));
    
    document.getElementById('bar-ollama').style.width = `${ollamaPct}%`;
    document.getElementById('bar-comfy').style.width = `${comfyPct}%`;
    document.getElementById('bar-system').style.width = `${systemPct}%`;
    document.getElementById('bar-free').style.width = `${freePct}%`;
    
    document.getElementById('tooltip-ollama').textContent = `${bd.ollama_gb || 0} GB`;
    document.getElementById('tooltip-comfy').textContent = `${bd.comfyui_gb || 0} GB`;
    document.getElementById('tooltip-system').textContent = `${bd.system_gb || 0} GB`;
    
    document.getElementById('legend-ollama').textContent = `${bd.ollama_gb || 0} GB`;
    document.getElementById('legend-comfy').textContent = `${bd.comfyui_gb || 0} GB`;
    document.getElementById('legend-system').textContent = `${bd.system_gb || 0} GB`;
    document.getElementById('legend-free').textContent = `${bd.free_gb || 0} GB`;
    
    // Hardware sensors
    document.getElementById('gpu-util-val').textContent = `${gpu.gpu_util_pct || 0}%`;
    document.getElementById('gpu-temp-val').textContent = `${gpu.temperature_c || 0}°C`;
    document.getElementById('gpu-power-val').textContent = `${gpu.power_w || 0} W`;
    document.getElementById('gpu-fan-val').textContent = `${gpu.fan_pct || 0}%`;
    
    // Processes table
    const tbody = document.getElementById('gpu-proc-table');
    if (bd.processes && bd.processes.length > 0) {
      tbody.innerHTML = bd.processes.map(p => {
        let tagClass = 'text-slate-400';
        let badge = '';
        if (p.is_ollama) {
          tagClass = 'text-purple-400 font-bold';
          badge = '<span class="px-1 py-0.2 bg-purple-950 text-purple-300 rounded border border-purple-800 text-[9px] mr-1">OLLAMA</span>';
        } else if (p.is_comfy) {
          tagClass = 'text-cyan-400 font-bold';
          badge = '<span class="px-1 py-0.2 bg-cyan-950 text-cyan-300 rounded border border-cyan-800 text-[9px] mr-1">COMFY</span>';
        }
        return `
          <tr class="hover:bg-slate-900/50">
            <td class="py-1 text-slate-500 font-mono text-[10px]">${p.pid}</td>
            <td class="py-1 ${tagClass} text-[11px] truncate max-w-[140px]">${badge}${p.name}</td>
            <td class="py-1 text-right font-mono font-bold text-slate-200 text-[11px]">${p.vram_mb} MB</td>
          </tr>
        `;
      }).join('');
    } else {
      tbody.innerHTML = '<tr><td colspan="3" class="py-2 text-center text-slate-500">No compute processes running</td></tr>';
    }
  }

  // 2. System RAM & Page Cache
  const ram = data.ram || {};
  if (ram.total_bytes) {
    document.getElementById('ram-cached-gb').textContent = ram.cached_gb || '0.0';
    document.getElementById('ram-total-text').textContent = `${ram.total_gb || 0} GB Total (${ram.cache_ratio_pct || 0}% in Cache)`;
    
    const usedPct = ((ram.used_bytes || 0) / ram.total_bytes) * 100;
    const cachePct = ((ram.cached_bytes || 0) / ram.total_bytes) * 100;
    const freePct = Math.max(0, 100 - (usedPct + cachePct));
    
    document.getElementById('bar-ram-used').style.width = `${usedPct}%`;
    document.getElementById('bar-ram-cache').style.width = `${cachePct}%`;
    document.getElementById('bar-ram-free').style.width = `${freePct}%`;
    
    document.getElementById('legend-ram-used').textContent = `${ram.used_gb || 0} GB`;
    document.getElementById('legend-ram-cached').textContent = `${ram.cached_gb || 0} GB`;
    document.getElementById('legend-ram-free').textContent = `${ram.free_gb || 0} GB`;
  }

  // 3. Ollama State
  const ollama = data.ollama || {};
  if (ollama.online) {
    document.getElementById('ollama-status-text').textContent = 'ONLINE';
    document.getElementById('ollama-status-text').className = 'text-emerald-400 font-mono font-bold';
    document.getElementById('ollama-dot').className = 'relative inline-flex rounded-full h-2.5 w-2.5 bg-emerald-500';
    document.getElementById('ollama-pulse').className = 'animate-ping absolute inline-flex h-full w-full rounded-full bg-emerald-400 opacity-75';
    
    if (ollama.active_model_name) {
      document.getElementById('ollama-active-model').textContent = ollama.active_model_name;
      document.getElementById('ollama-vram-badge').textContent = `${ollama.active_model_vram_gb || 0} GB VRAM`;
      document.getElementById('ollama-vram-badge').className = 'px-2 py-0.5 text-[10px] font-mono rounded bg-purple-900/80 text-purple-200 border border-purple-600 font-bold';
      document.getElementById('ollama-context').textContent = `${ollama.active_context || 0} ctx`;
    } else {
      document.getElementById('ollama-active-model').textContent = 'None Loaded (VRAM Free)';
      document.getElementById('ollama-vram-badge').textContent = '0.0 GB VRAM';
      document.getElementById('ollama-vram-badge').className = 'px-2 py-0.5 text-[10px] font-mono rounded bg-slate-800 text-slate-400 border border-slate-700';
      document.getElementById('ollama-context').textContent = 'Idle';
    }
    
    if (ollama.installed_models && ollama.installed_models.length > 0) {
      document.getElementById('ollama-total-models').textContent = ollama.installed_models.length;
      updateModelSelect(ollama.installed_models, ollama.active_model_name);
    }
  } else {
    document.getElementById('ollama-status-text').textContent = 'OFFLINE';
    document.getElementById('ollama-status-text').className = 'text-rose-400 font-mono font-bold';
    document.getElementById('ollama-dot').className = 'relative inline-flex rounded-full h-2.5 w-2.5 bg-rose-500';
    document.getElementById('ollama-pulse').className = 'hidden';
  }

  // 4. ComfyUI State
  const comfy = data.comfyui || {};
  if (comfy.online) {
    document.getElementById('comfy-status-text').textContent = 'ONLINE';
    document.getElementById('comfy-status-text').className = 'text-emerald-400 font-mono font-bold';
    document.getElementById('comfy-dot').className = 'relative inline-flex rounded-full h-2.5 w-2.5 bg-emerald-500';
    document.getElementById('comfy-pulse').className = 'animate-ping absolute inline-flex h-full w-full rounded-full bg-emerald-400 opacity-75';
    
    document.getElementById('comfy-queue-count').textContent = comfy.queue_running + comfy.queue_remaining;
    if (comfy.executing) {
      document.getElementById('comfy-exec-badge').textContent = 'GENERATING...';
      document.getElementById('comfy-exec-badge').className = 'px-2 py-0.5 text-[10px] font-mono rounded bg-amber-950 text-amber-300 border border-amber-600 animate-pulse font-bold';
      document.getElementById('comfy-model-info').textContent = `Prompt #${comfy.current_prompt_id || 'Active'}`;
    } else {
      document.getElementById('comfy-exec-badge').textContent = 'IDLE / READY';
      document.getElementById('comfy-exec-badge').className = 'px-2 py-0.5 text-[10px] font-mono rounded bg-emerald-950 text-emerald-300 border border-emerald-800';
      document.getElementById('comfy-model-info').textContent = 'Dynamic Model Offloader';
    }
    
    document.getElementById('comfy-vram-avail').textContent = `${(comfy.vram_free_mb / 1024).toFixed(1)} GB`;
  } else {
    document.getElementById('comfy-status-text').textContent = 'OFFLINE';
    document.getElementById('comfy-status-text').className = 'text-rose-400 font-mono font-bold';
    document.getElementById('comfy-dot').className = 'relative inline-flex rounded-full h-2.5 w-2.5 bg-rose-500';
    document.getElementById('comfy-pulse').className = 'hidden';
  }
  
  document.getElementById('comfy-discovered-count').textContent = `${data.comfy_models_count || 0} Files`;

  // 5. Switch History Timeline
  const history = data.history || [];
  const logContainer = document.getElementById('switch-log-container');
  if (history.length > 0) {
    const latest = history[0];
    document.getElementById('ollama-last-swap').textContent = `${latest.duration_ms} ms`;
    document.getElementById('ollama-cache-hit').textContent = latest.cache_status || 'OK';
    
    logContainer.innerHTML = history.slice(0, 10).map(item => {
      const isHit = (item.cache_status || '').includes('RAM Cache Hit') || (item.cache_status || '').includes('RAM-Cached');
      const badgeClass = isHit 
        ? 'bg-emerald-950/80 text-emerald-300 border-emerald-800' 
        : 'bg-amber-950/80 text-amber-300 border-amber-800';
      return `
        <div class="p-2.5 rounded-xl bg-slate-950 border border-slate-800 text-xs font-mono flex items-center justify-between">
          <div class="space-y-0.5">
            <div class="flex items-center space-x-1.5">
              <span class="text-slate-500 text-[10px]">${item.timestamp}</span>
              <span class="text-indigo-400 font-bold">${item.event_type}</span>
            </div>
            <div class="text-[11px] text-slate-300 truncate max-w-[260px]">
              <span class="text-slate-500">${item.source}</span> → <span class="text-cyan-300 font-bold">${item.target}</span>
            </div>
          </div>
          <div class="text-right">
            <div class="text-slate-100 font-bold">${item.duration_ms} ms</div>
            <span class="px-1.5 py-0.2 text-[9px] rounded border ${badgeClass}">${item.cache_status}</span>
          </div>
        </div>
      `;
    }).join('');
  }
}

function updateModelSelect(models, activeModel) {
  const select = document.getElementById('ollama-model-select');
  const currentVal = select.value;
  
  if (JSON.stringify(models.map(m => m.name)) === JSON.stringify(currentInstalledModels)) {
    return;
  }
  
  currentInstalledModels = models.map(m => m.name);
  select.innerHTML = models.map(m => {
    const isSelected = m.name === activeModel || m.name === currentVal;
    const sizeGb = (m.size / (1024**3)).toFixed(1);
    const quant = m.details?.quantization_level || '';
    return `<option value="${m.name}" ${isSelected ? 'selected' : ''}>${m.name} (${sizeGb} GB ${quant})</option>`;
  }).join('');
}

// User Actions
async function triggerModelSwitch() {
  const select = document.getElementById('ollama-model-select');
  const targetModel = select.value;
  if (!targetModel) return;
  
  const btn = document.getElementById('btn-switch-model');
  btn.disabled = true;
  btn.innerHTML = '<i class="fa-solid fa-spinner fa-spin"></i> <span>Swapping...</span>';
  
  try {
    const resp = await fetch('/api/switch-model', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ model: targetModel, keep_alive: '30m' })
    });
    const result = await resp.json();
    if (!resp.ok) {
      alert(`Switch failed: ${result.detail || 'Unknown error'}`);
    }
  } catch (err) {
    alert(`Error: ${err}`);
  } finally {
    btn.disabled = false;
    btn.innerHTML = '<i class="fa-solid fa-shuffle"></i> <span>Hot Swap</span>';
  }
}

async function freeOllamaVRAM() {
  try {
    const resp = await fetch('/api/free-vram', { method: 'POST' });
    const data = await resp.json();
    console.log('Ollama VRAM yielded:', data);
  } catch (err) {
    alert(`Error: ${err}`);
  }
}

async function freeComfyVRAM() {
  try {
    const resp = await fetch('/api/comfy-free', { method: 'POST' });
    const data = await resp.json();
    console.log('ComfyUI VRAM freed:', data);
  } catch (err) {
    alert(`Error: ${err}`);
  }
}

async function warmAllModels() {
  try {
    const resp = await fetch('/api/warm-all', { method: 'POST' });
    const data = await resp.json();
    alert(`Warming completed in ${data.total_duration_ms} ms! All models are now cached in 64GB RAM.`);
  } catch (err) {
    alert(`Error warming models: ${err}`);
  }
}

// Startup
document.addEventListener('DOMContentLoaded', () => {
  initSSE();
});