// 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
  services: {},
  service_action_in_flight: null,  // e.g. "parakeet:restart"
  configured: true,
  timer_handle: null,
};

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

function escapeHtml(s) {
  if (s == null) return '';
  return String(s)
    .replaceAll('&', '&')
    .replaceAll('<', '&lt;')
    .replaceAll('>', '&gt;')
    .replaceAll('"', '&quot;')
    .replaceAll("'", '&#39;');
}

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' : '');
    const desc = m.description
      ? `<div class="desc">${escapeHtml(m.description)}</div>`
      : '';
    card.innerHTML = `
      <div class="name">${escapeHtml(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">${escapeHtml(c)}</span>`).join('')}
      </div>
      ${desc}
      <div class="muted small repo">${escapeHtml(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 classifyService(s) {
  // returns one of: running | unhealthy | missing | unconfigured | starting
  if (!s.host) return 'unconfigured';
  if (s.docker_state === 'missing') return 'missing';
  if (s.docker_state === 'restarting') return 'unhealthy';
  if (s.docker_state === 'exited') return 'unhealthy';
  if (s.docker_state === 'running' && !s.http_ready) return 'starting';
  if (s.docker_state === 'running' && s.http_ready) return 'running';
  return s.docker_state || 'unknown';
}

function statusLabel(cls) {
  return {
    running: 'Healthy',
    unhealthy: 'Unhealthy',
    starting: 'Starting…',
    missing: 'Not installed',
    unconfigured: 'Not configured',
    unknown: 'Unknown',
  }[cls] || cls;
}

async function renderServices() {
  let services = state.services;
  // First render: fetch.
  if (!services || Object.keys(services).length === 0) {
    try {
      services = await fetchJSON('/api/services');
      state.services = services;
    } catch (e) { console.error('services fetch failed', e); return; }
  }
  const panel = el('#services-panel');
  const grid = el('#services-grid');
  const entries = Object.entries(services);
  if (entries.length === 0) { panel.classList.add('hidden'); return; }
  panel.classList.remove('hidden');
  grid.innerHTML = '';
  for (const [name, s] of entries) {
    const cls = classifyService(s);
    const card = document.createElement('div');
    card.className = `service-card ${cls}`;
    const inFlight = state.service_action_in_flight && state.service_action_in_flight.startsWith(name + ':');
    const disable = (action) => {
      // Disable buttons that don't make sense for the current state
      if (inFlight) return true;
      if (cls === 'unconfigured' || cls === 'missing') return true;
      if (action === 'start' && (cls === 'running' || cls === 'starting')) return true;
      if (action === 'stop' && cls !== 'running' && cls !== 'starting' && cls !== 'unhealthy') return true;
      return false;
    };
    const hostRow = s.host
      ? `<div class="row"><span class="k">Host</span><span class="v">${escapeHtml(s.host)}:${s.port}</span></div>`
      : `<div class="row"><span class="k">Host</span><span class="v muted-v">not configured</span></div>`;
    const modelRow = s.model
      ? `<div class="row"><span class="k">Model</span><span class="v">${escapeHtml(s.model)}</span></div>`
      : '';
    const restartsRow = s.restart_count != null && s.restart_count > 1
      ? `<div class="row"><span class="k">Restarts</span><span class="v">${s.restart_count}</span></div>`
      : '';
    card.innerHTML = `
      <div class="head">
        <span class="name">${escapeHtml(name)}</span>
        <span class="kind">${escapeHtml(s.kind || '')}</span>
        <span class="status">${statusLabel(cls)}</span>
      </div>
      ${hostRow}
      ${modelRow}
      ${restartsRow}
      <div class="service-actions">
        <button class="btn" data-svc-action="${name}:start" ${disable('start') ? 'disabled' : ''}>Start</button>
        <button class="btn" data-svc-action="${name}:restart" ${disable('restart') ? 'disabled' : ''}>Restart</button>
        <button class="btn danger" data-svc-action="${name}:stop" ${disable('stop') ? 'disabled' : ''}>Stop</button>
      </div>
    `;
    grid.appendChild(card);
  }
  for (const btn of grid.querySelectorAll('.btn[data-svc-action]')) {
    btn.addEventListener('click', () => onServiceAction(btn.dataset.svcAction));
  }
}

async function onServiceAction(key) {
  if (state.service_action_in_flight) return;
  const [name, action] = key.split(':');
  state.service_action_in_flight = key;
  renderServices();
  try {
    await fetchJSON(`/api/services/${name}/${action}`, { method: 'POST' });
  } catch (e) {
    alert(`${action} ${name} failed: ${e.message}`);
  } finally {
    state.service_action_in_flight = null;
    // Refresh services state
    try {
      state.services = await fetchJSON('/api/services');
    } catch {}
    renderServices();
    pollStatus();
  }
}

function renderEndpoint(status) {
  const v = status.vllm || {};
  const panel = el('#endpoint-panel');
  const ready = v.ok && v.current_model && v.base_url;
  panel.classList.toggle('hidden', !ready);
  if (!ready) return;
  el('#ep-url').textContent = v.base_url;
  el('#ep-model').textContent = v.current_model;
  const snippet =
`curl -s ${v.base_url}/chat/completions \\
  -H 'content-type: application/json' \\
  -d '{
    "model": "${v.current_model}",
    "messages": [{"role": "user", "content": "Hello"}]
  }'`;
  el('#ep-curl-snippet').textContent = snippet;
}

function setupCopyButtons() {
  document.body.addEventListener('click', async (e) => {
    const btn = e.target.closest('.copy-btn');
    if (!btn) return;
    const targetSel = btn.dataset.copy;
    if (!targetSel) return;
    const target = el(targetSel);
    if (!target) return;
    const text = target.textContent;
    try {
      await navigator.clipboard.writeText(text);
      const original = btn.textContent;
      btn.classList.add('copied');
      btn.textContent = 'Copied';
      setTimeout(() => {
        btn.classList.remove('copied');
        btn.textContent = original;
      }, 1400);
    } catch {
      // Clipboard API may fail over plain HTTP; fall back to selection
      const range = document.createRange();
      range.selectNode(target);
      window.getSelection().removeAllRanges();
      window.getSelection().addRange(range);
    }
  });
}

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);
    renderEndpoint(status);
    renderHealth(status);
    // Refresh services state lazily — every 5s poll triggers this too.
    try {
      state.services = await fetchJSON('/api/services');
      renderServices();
    } catch {}
    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() {
  setupCopyButtons();
  await loadModels();
  await pollStatus();
  await renderServices();
  setInterval(pollStatus, 5000);
}

init();