// spark-control front-end

const state = {
  models: {},
  defaults: {},
  current_model_key: null,
  swap_job_id: null,
  swap_eventsource: null,
  swap_started_at: null,
  swap_lines: [],          // local accumulator for phase detection
  swap_phase: 'Starting…',
  swap_phase_detail: '',
  swap_progress: 0,        // 0–1
  configured: true,
  timer_handle: null,
};

const el = (sel) => document.querySelector(sel);
const $$ = (sel) => document.querySelectorAll(sel);

async function fetchJSON(url, opts) {
  const r = await fetch(url, opts);
  if (!r.ok) {
    const text = await r.text().catch(() => '');
    throw new Error(`${r.status} ${r.statusText}: ${text}`);
  }
  return r.json();
}

// ===================== rendering =====================

function renderCards() {
  const root = el('#cards');
  root.innerHTML = '';
  const isSwapping = !!state.swap_job_id;
  for (const key of Object.keys(state.models)) {
    const m = state.models[key];
    const isActive = key === state.current_model_key;
    const card = document.createElement('div');
    card.className = 'card' + (isActive ? ' active' : '');
    card.innerHTML = `
      <div class="name">${m.display_name}</div>
      <div class="meta">
        <span class="tag mode-${m.mode}">${m.mode}</span>
        <span class="tag">${m.size_gb} GB</span>
        ${(m.capabilities || []).map(c => `<span class="tag cap">${c}</span>`).join('')}
      </div>
      <div class="muted small repo">${m.repo}</div>
      <div class="spacer"></div>
      <button class="btn ${isActive ? '' : 'primary'}" data-key="${key}" ${isActive || isSwapping ? 'disabled' : ''}>
        ${isActive ? 'Current' : 'Switch to this'}
      </button>
    `;
    root.appendChild(card);
  }
  for (const btn of $$('.card .btn')) {
    btn.addEventListener('click', () => triggerSwap(btn.dataset.key));
  }
}

function renderCurrent(status) {
  const c = el('#current');
  if (!status.configured) { c.innerHTML = `<span class="muted">not configured</span>`; return; }
  if (status.current_swap_job) { c.innerHTML = `<span class="muted">swap in progress</span>`; return; }
  const v = status.vllm || {};
  if (!v.ok) { c.innerHTML = `<span class="muted">vLLM unreachable</span>`; return; }
  const m = status.current_model_key ? state.models[status.current_model_key] : null;
  const label = m ? m.display_name : (v.current_model || '(unknown)');
  c.innerHTML = `<strong>${label}</strong>`;
}

function renderHealth(status) {
  function setDot(id, ok, payload) {
    const item = el(id);
    if (!item) return;
    const dot = item.querySelector('.dot');
    dot.classList.remove('ok', 'bad', 'warn');
    if (ok === true) dot.classList.add('ok');
    else if (ok === false) dot.classList.add('bad');
    else dot.classList.add('warn');
    item.title = JSON.stringify(payload || {}, null, 2);
  }
  setDot('#h-vllm', status.vllm && status.vllm.ok, status.vllm);
  setDot('#h-parakeet', status.parakeet && status.parakeet.ok, status.parakeet);
  setDot('#h-magpie', status.magpie && status.magpie.ok, status.magpie);
  el('#updated').textContent = `updated ${new Date().toLocaleTimeString()}`;
}

function renderBanner(status) {
  el('#setup-banner').classList.toggle('hidden', !!status.configured);
}

function renderSwapPanel() {
  el('#swap-phase').textContent = state.swap_phase;
  el('#swap-phase-detail').textContent = state.swap_phase_detail;
  el('#swap-phase-fill').style.width = `${Math.max(2, Math.round(state.swap_progress * 100))}%`;
}

// ===================== phase detection =====================

const PHASE_ORDER = [
  ['Stopping current model…', 0.08],
  ['Starting new model…',     0.16],
  ['Joining Ray cluster…',    0.22],
  ['Loading weights…',        0.30],
  ['Compiling kernels…',      0.78],
  ['Warming up…',             0.88],
  ['Starting API server…',    0.94],
  ['Ready ✓',                 1.00],
  ['Failed',                  1.00],
];

function phaseProgress(name) {
  const found = PHASE_ORDER.find(([n]) => n === name);
  return found ? found[1] : 0.05;
}

function deriveSwapPhase(serverState, lines) {
  // Default phase from server state
  let phase = ({
    starting:  'Starting…',
    stopping:  'Stopping current model…',
    launching: 'Starting new model…',
    tailing:   'Loading weights…',
    ready:     'Ready ✓',
    failed:    'Failed',
  })[serverState] || 'Working…';

  let detail = '';

  // Refine from log content (search recent lines first)
  const tail = lines.slice(-40);
  for (let i = tail.length - 1; i >= 0; i--) {
    const line = tail[i];
    if (line.includes('Application startup complete')) {
      phase = 'Ready ✓';
      break;
    }
    if (line.includes('Started server process')) {
      phase = 'Starting API server…';
      break;
    }
    if (line.includes('Initial profiling/warmup') || line.includes('init engine (profile, create kv cache, warmup model)')) {
      phase = 'Warming up…';
      break;
    }
    if (line.match(/Capturing CUDA graphs|Compiling a graph|torch\.compile took|Graph capturing/)) {
      phase = 'Compiling kernels…';
      break;
    }
    const shard = line.match(/Loading safetensors checkpoint shards:\s+(\d+)%\s+Completed\s+\|\s+(\d+)\/(\d+)/);
    if (shard) {
      phase = 'Loading weights…';
      detail = `${shard[2]} of ${shard[3]} shards (${shard[1]}%)`;
      const innerProgress = parseInt(shard[2], 10) / parseInt(shard[3], 10);
      // Map shard progress 0..1 into the 0.30..0.78 band
      state.swap_progress = 0.30 + (0.78 - 0.30) * innerProgress;
      state.swap_phase = phase;
      state.swap_phase_detail = detail;
      return;
    }
    if (line.includes('Connecting to existing Ray cluster')) {
      phase = 'Joining Ray cluster…';
      break;
    }
    if (line.includes('Resolved architecture') || line.match(/launch-cluster\.sh.*exec vllm serve/)) {
      phase = 'Starting new model…';
      break;
    }
    if (line.match(/launch-cluster\.sh stop/)) {
      phase = 'Stopping current model…';
      break;
    }
  }

  state.swap_phase = phase;
  state.swap_phase_detail = detail;
  state.swap_progress = phaseProgress(phase);
}

// ===================== timer =====================

function startTimer(startedAtMillis) {
  state.swap_started_at = startedAtMillis;
  if (state.timer_handle) clearInterval(state.timer_handle);
  const tick = () => {
    if (!state.swap_started_at) return;
    const sec = Math.max(0, Math.floor((Date.now() - state.swap_started_at) / 1000));
    const m = Math.floor(sec / 60);
    const s = sec % 60;
    el('#swap-elapsed').textContent = `${m}:${s.toString().padStart(2, '0')}`;
  };
  tick();
  state.timer_handle = setInterval(tick, 500);
}

function stopTimer() {
  if (state.timer_handle) { clearInterval(state.timer_handle); state.timer_handle = null; }
}

// ===================== polling + SSE =====================

async function pollStatus() {
  try {
    const status = await fetchJSON('/api/status');
    state.current_model_key = status.current_model_key;
    state.configured = status.configured;
    renderBanner(status);
    renderCurrent(status);
    renderHealth(status);
    if (status.current_swap_job && status.current_swap_job !== state.swap_job_id) {
      attachToSwap(status.current_swap_job, /*needsBackfill=*/true);
    } else if (!status.current_swap_job && state.swap_job_id && !state.swap_eventsource) {
      // Foreign swap ended
      detachSwap();
    }
    renderCards();
  } catch (e) {
    console.error('status poll failed', e);
  }
}

async function loadModels() {
  const data = await fetchJSON('/api/models');
  state.defaults = data.defaults || {};
  state.models = data.models || {};
}

async function triggerSwap(modelKey) {
  if (state.swap_job_id) return;
  try {
    const r = await fetchJSON('/api/swap', {
      method: 'POST',
      headers: { 'content-type': 'application/json' },
      body: JSON.stringify({ model_key: modelKey }),
    });
    attachToSwap(r.job_id, /*needsBackfill=*/false);
  } catch (e) {
    alert('Failed to start swap: ' + e.message);
  }
}

async function attachToSwap(jobId, needsBackfill) {
  if (state.swap_eventsource) {
    state.swap_eventsource.close();
    state.swap_eventsource = null;
  }
  state.swap_job_id = jobId;
  state.swap_lines = [];
  state.swap_phase = 'Starting…';
  state.swap_phase_detail = '';
  state.swap_progress = 0.05;
  el('#swap-log').textContent = '';
  el('#swap-panel').classList.remove('hidden');
  renderSwapPanel();

  // Backfill (if joining mid-swap) — fetch the snapshot so we have started_at + history
  try {
    const snap = await fetchJSON(`/api/swap/${jobId}`);
    const ts = Date.parse(snap.started_at);
    if (!isNaN(ts)) startTimer(ts);
    state.swap_lines = snap.lines || [];
    for (const line of state.swap_lines) appendLog(line);
    deriveSwapPhase(snap.state, state.swap_lines);
    renderSwapPanel();
    if (snap.returncode !== null && snap.returncode !== undefined) {
      // Already finished — close panel after a beat
      handleSwapDone(snap);
      return;
    }
  } catch (e) {
    if (!needsBackfill) startTimer(Date.now());
    console.warn('backfill failed', e);
  }

  const es = new EventSource(`/api/swap/${jobId}/stream`);
  state.swap_eventsource = es;
  es.onmessage = (ev) => {
    try {
      const d = JSON.parse(ev.data);
      if (d.line !== undefined) {
        state.swap_lines.push(d.line);
        appendLog(d.line);
        deriveSwapPhase(d.state, state.swap_lines);
        renderSwapPanel();
      } else if (d.state) {
        deriveSwapPhase(d.state, state.swap_lines);
        renderSwapPanel();
      }
    } catch {}
  };
  es.addEventListener('done', async (ev) => {
    let d = {};
    try { d = JSON.parse(ev.data); } catch {}
    handleSwapDone(d);
  });
  es.onerror = () => {
    // Tab backgrounded or network blip — close; status poll will reattach
    es.close();
    state.swap_eventsource = null;
  };
  renderCards();
}

function handleSwapDone(d) {
  if (state.swap_eventsource) { state.swap_eventsource.close(); state.swap_eventsource = null; }
  const finalState = d.state || 'ready';
  state.swap_phase = finalState === 'failed' ? 'Failed' : 'Ready ✓';
  state.swap_phase_detail = d.returncode !== undefined ? `exit code ${d.returncode}` : '';
  state.swap_progress = 1.0;
  renderSwapPanel();
  setTimeout(() => detachSwap(), 4000);
  pollStatus();
}

function detachSwap() {
  state.swap_job_id = null;
  if (state.swap_eventsource) { state.swap_eventsource.close(); state.swap_eventsource = null; }
  stopTimer();
  el('#swap-panel').classList.add('hidden');
  renderCards();
}

function appendLog(line) {
  const log = el('#swap-log');
  log.textContent += line + '\n';
  log.scrollTop = log.scrollHeight;
}

async function init() {
  await loadModels();
  await pollStatus();
  setInterval(pollStatus, 5000);
}

init();