/* ─── Config ──────────────────────────────────────────────────────────────── */
// Empty string = same origin (works on HF Spaces and any deployment).
// Falls back to localhost only when running locally on a different port.
const API_BASE = window.location.hostname === 'localhost' ? 'http://localhost:8000' : '';

/* ─── State ───────────────────────────────────────────────────────────────── */
const state = {
  currentTaskId: null,
  isRunning: false,
  taskHistory: JSON.parse(localStorage.getItem('taskHistory') || '[]'),
};

/* ─── DOM refs ────────────────────────────────────────────────────────────── */
const $ = id => document.getElementById(id);
const taskInput    = $('taskInput');
const submitBtn    = $('submitBtn');
const submitText   = $('submitText');
const charCount    = $('charCount');
const eventLog     = $('eventLog');
const planCard     = $('planCard');
const planSteps    = $('planSteps');
const planProgress = $('planProgress');
const outputCard   = $('outputCard');
const outputContent = $('outputContent');
const qualityBadge = $('qualityBadge');
const taskHistory  = $('taskHistory');
const globalStatus  = $('globalStatus');
const healthDot     = $('healthDot');
const healthLabel   = $('healthLabel');
const errorBanner   = $('errorBanner');
const errorBannerMsg = $('errorBannerMsg');
$('errorBannerClose').addEventListener('click', () => { errorBanner.style.display = 'none'; });

/* ─── Health check ────────────────────────────────────────────────────────── */
async function checkHealth() {
  try {
    const res = await fetch(`${API_BASE}/api/health`, { signal: AbortSignal.timeout(3000) });
    if (res.ok) {
      healthDot.className = 'health-indicator ok';
      healthLabel.textContent = 'Backend online';
    } else {
      throw new Error('non-ok');
    }
  } catch {
    healthDot.className = 'health-indicator error';
    healthLabel.textContent = 'Backend offline';
  }
}
checkHealth();
setInterval(checkHealth, 15000);

/* ─── Theme toggle ────────────────────────────────────────────────────────── */
const themeToggle = $('themeToggle');
let isDark = true;
themeToggle.addEventListener('click', () => {
  isDark = !isDark;
  document.documentElement.setAttribute('data-theme', isDark ? '' : 'light');
  themeToggle.textContent = isDark ? '☀' : '🌙';
});

/* ─── Char counter ────────────────────────────────────────────────────────── */
taskInput.addEventListener('input', () => {
  const len = taskInput.value.length;
  charCount.textContent = `${len} / 2000`;
  charCount.style.color = len > 1800 ? 'var(--amber)' : '';
});

/* ─── Sample chips ────────────────────────────────────────────────────────── */
document.querySelectorAll('.chip').forEach(chip => {
  chip.addEventListener('click', () => {
    taskInput.value = chip.textContent.trim();
    taskInput.dispatchEvent(new Event('input'));
    taskInput.focus();
  });
});

/* ─── Clear events ────────────────────────────────────────────────────────── */
$('clearEvents').addEventListener('click', () => {
  eventLog.innerHTML = '<p class="empty-state">Events will appear here</p>';
});

/* ─── Copy output ─────────────────────────────────────────────────────────── */
$('copyBtn').addEventListener('click', () => {
  navigator.clipboard.writeText(outputContent.textContent).then(() => {
    const btn = $('copyBtn');
    btn.textContent = 'Copied!';
    setTimeout(() => { btn.textContent = 'Copy'; }, 1500);
  });
});

/* ─── Agent state management ─────────────────────────────────────────────── */
const AGENT_ORDER = ['memory_retrieve', 'planner', 'executor', 'critic', 'memory_store'];

const AGENT_ROLE_MAP = {
  memory: 'memory_retrieve',
  planner: 'planner',
  executor: 'executor',
  critic: 'critic',
  memory_store: 'memory_store',
};

function setAgentState(agentId, stateStr) {
  const node = $(`node-${agentId}`);
  if (!node) return;
  node.className = `agent-node ${stateStr}`;
  const badge = node.querySelector('.node-state');
  badge.className = `node-state ${stateStr}`;
  badge.textContent = stateStr;
}

function resetAgentNodes() {
  AGENT_ORDER.forEach(id => setAgentState(id, 'idle'));
}

/* ─── Event log ───────────────────────────────────────────────────────────── */
function addEvent(agent, eventType, message) {
  const empty = eventLog.querySelector('.empty-state');
  if (empty) empty.remove();

  const time = new Date().toLocaleTimeString([], { hour: '2-digit', minute: '2-digit', second: '2-digit' });
  const div = document.createElement('div');

  const roleClass = (agent || '').toLowerCase().replace(' ', '_');
  div.className = `event-item ${roleClass}`;
  div.innerHTML = `
    <div class="event-header">
      <span class="event-agent">${agent || 'system'} · ${eventType || ''}</span>
      <span class="event-time">${time}</span>
    </div>
    <div class="event-msg">${escHtml(message || '')}</div>
  `;
  eventLog.prepend(div);

  // Cap at 50 events
  while (eventLog.children.length > 50) {
    eventLog.removeChild(eventLog.lastChild);
  }
}

function escHtml(str) {
  return str.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;');
}

/* ─── Plan rendering ─────────────────────────────────────────────────────── */
function renderPlan(plan) {
  if (!plan || !plan.length) return;
  planCard.style.display = '';

  const done = plan.filter(s => s.status === 'done').length;
  planProgress.textContent = `${done} / ${plan.length} steps`;

  planSteps.innerHTML = plan.map((step, i) => {
    const status = step.status || 'pending';
    const numHtml = status === 'done'
      ? '✓'
      : status === 'failed' ? '✗' : i + 1;
    return `
      <div class="step-item ${status}">
        <div class="step-num">${numHtml}</div>
        <div class="step-body">
          <div class="step-title">${escHtml(step.title || '')}</div>
          <div class="step-desc">${escHtml(step.description || '')}</div>
          ${step.tool ? `<span class="step-tool">${escHtml(step.tool)}</span>` : ''}
        </div>
      </div>`;
  }).join('');
}

/* ─── Metrics update ─────────────────────────────────────────────────────── */
function updateMetrics({ total_tokens, plan, iteration, quality_score }) {
  if (total_tokens != null) $('metricTokens').textContent = total_tokens.toLocaleString();
  if (plan != null)         $('metricSteps').textContent = plan.length;
  if (iteration != null)    $('metricIter').textContent = iteration;
  if (quality_score != null) {
    const s = Math.round(quality_score);
    $('metricScore').textContent = s;
    qualityBadge.textContent = `${s}/100`;
    qualityBadge.style.background = s >= 80 ? 'rgba(16,185,129,0.2)' : s >= 60 ? 'rgba(245,158,11,0.2)' : 'rgba(239,68,68,0.2)';
    qualityBadge.style.color = s >= 80 ? 'var(--green)' : s >= 60 ? 'var(--amber)' : 'var(--red)';
  }
}

/* ─── Error banner ───────────────────────────────────────────────────────── */
function showError(msg) {
  errorBannerMsg.textContent = msg;
  errorBanner.style.display = 'flex';
  addEvent('system', 'error', msg);
}

/* ─── Status badge ───────────────────────────────────────────────────────── */
function setStatus(status) {
  globalStatus.textContent = status;
  globalStatus.className = `status-badge ${(status || '').toLowerCase().replace(/[^a-z]/g, '')}`;
}

/* ─── Map workflow status → active agent node ───────────────────────────── */
function syncAgentsToStatus(status, latestEvent) {
  const role = latestEvent?.agent_role?.toLowerCase() || '';

  // Determine which agent is currently active
  const active = role.includes('memory') && status === 'planning'
    ? 'memory_retrieve'
    : role.includes('memory')
      ? (status === 'completed' ? 'memory_store' : 'memory_retrieve')
    : role.includes('planner') ? 'planner'
    : role.includes('executor') ? 'executor'
    : role.includes('critic') ? 'critic'
    : null;

  if (!active) return;

  // Mark preceding agents as done, active as running, rest idle
  let passed = false;
  AGENT_ORDER.forEach(id => {
    if (id === active) {
      setAgentState(id, 'running');
      passed = true;
    } else if (!passed) {
      setAgentState(id, 'done');
    } else {
      setAgentState(id, 'idle');
    }
  });
}

/* ─── Mark all done ──────────────────────────────────────────────────────── */
function markAllDone() {
  AGENT_ORDER.forEach(id => setAgentState(id, 'done'));
}

/* ─── Submit task ────────────────────────────────────────────────────────── */
submitBtn.addEventListener('click', runTask);
taskInput.addEventListener('keydown', e => {
  if (e.key === 'Enter' && (e.ctrlKey || e.metaKey)) runTask();
});

async function runTask() {
  const task = taskInput.value.trim();
  if (!task || state.isRunning) return;

  state.isRunning = true;
  submitBtn.disabled = true;
  submitText.innerHTML = '<span class="spinner"></span> Running…';
  setStatus('planning');
  resetAgentNodes();
  planCard.style.display = 'none';
  outputCard.style.display = 'none';
  planSteps.innerHTML = '';
  $('metricTokens').textContent = '—';
  $('metricSteps').textContent = '—';
  $('metricIter').textContent = '—';
  $('metricScore').textContent = '—';
  qualityBadge.textContent = '';

  errorBanner.style.display = 'none';
  addEvent('system', 'task_submitted', task.slice(0, 120));

  const useStream = $('streamToggle').checked;

  try {
    if (useStream) {
      await runStreaming(task);
    } else {
      await runBatch(task);
    }
  } finally {
    state.isRunning = false;
    submitBtn.disabled = false;
    submitText.innerHTML = 'Run Agents <svg class="btn-icon-arrow" width="16" height="16" viewBox="0 0 16 16" fill="none"><path d="M3 8h10M9 4l4 4-4 4" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"/></svg>';
  }
}

/* ─── Streaming mode ─────────────────────────────────────────────────────── */
async function runStreaming(task) {
  const res = await fetch(`${API_BASE}/api/tasks/stream`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ task, stream: true }),
  });

  if (!res.ok) {
    addEvent('system', 'error', `HTTP ${res.status}: ${await res.text()}`);
    setStatus('failed');
    return;
  }

  const reader = res.body.getReader();
  const decoder = new TextDecoder();
  let buffer = '';

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    buffer += decoder.decode(value, { stream: true });

    const lines = buffer.split('\n');
    buffer = lines.pop();

    let pendingEventType = null;
    for (const line of lines) {
      if (line.startsWith('event: ')) {
        pendingEventType = line.slice(7).trim();
      } else if (line.startsWith('data: ')) {
        try {
          const data = JSON.parse(line.slice(6));
          if (pendingEventType === 'done') {
            markAllDone();
            setStatus('completed');
          } else if (pendingEventType === 'error') {
            setStatus('failed');
            showError(data.error || 'Task failed — check server logs for details.');
          } else {
            handleStreamUpdate(data);
          }
        } catch { /* ignore parse errors */ }
        pendingEventType = null;
      }
    }
  }
}

function handleStreamUpdate(data) {
  const { status, plan, latest_event, quality_score, total_tokens, iteration, task_id, final_output, error_message } = data;

  if (task_id) state.currentTaskId = task_id;

  if (status) setStatus(status);

  if (latest_event?.agent_role) {
    syncAgentsToStatus(status, latest_event);
    addEvent(latest_event.agent_role, latest_event.event_type, latest_event.message);
  }

  // Surface agent-level errors even when sent as a state update
  if (latest_event?.event_type === 'error') {
    showError(latest_event.message || error_message || 'Agent error');
  }

  if (plan?.length) renderPlan(plan);

  updateMetrics({ total_tokens, plan, iteration, quality_score });

  if (final_output) showOutput(final_output, quality_score);

  if (status === 'completed') {
    markAllDone();
    saveToHistory({ task: taskInput.value.trim(), task_id, status: 'completed', score: quality_score });
  } else if (status === 'failed') {
    const errMsg = error_message || latest_event?.message || 'Task failed';
    showError(errMsg);
    saveToHistory({ task: taskInput.value.trim(), task_id, status: 'failed', score: null });
  }
}

/* ─── Batch mode ─────────────────────────────────────────────────────────── */
async function runBatch(task) {
  setAgentState('memory_retrieve', 'running');

  const res = await fetch(`${API_BASE}/api/tasks`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ task }),
  });

  if (!res.ok) {
    const err = await res.text();
    addEvent('system', 'error', `HTTP ${res.status}: ${err}`);
    setStatus('failed');
    return;
  }

  const data = await res.json();
  state.currentTaskId = data.task_id;

  // Replay events
  (data.events || []).forEach(ev => {
    addEvent(ev.agent_role, ev.event_type, ev.message);
  });

  setStatus(data.status);
  renderPlan(data.plan);
  updateMetrics({
    total_tokens: data.total_tokens,
    plan: data.plan,
    iteration: null,
    quality_score: data.quality_score,
  });

  if (data.final_output) showOutput(data.final_output, data.quality_score);

  if (data.status === 'completed') {
    markAllDone();
    saveToHistory({ task, task_id: data.task_id, status: 'completed', score: data.quality_score });
  } else {
    const errMsg = data.error_message || 'Task failed — check server logs.';
    showError(errMsg);
    saveToHistory({ task, task_id: data.task_id, status: data.status, score: null });
  }
}

/* ─── Show output ────────────────────────────────────────────────────────── */
function cleanOutput(text) {
  // Strip LaTeX boxing that Llama sometimes outputs
  return (text || '').replace(/\$\\boxed\{([^}]+)\}\$/g, '$1').replace(/\\boxed\{([^}]+)\}/g, '$1');
}

function showOutput(text, score) {
  outputCard.style.display = '';
  outputContent.textContent = cleanOutput(text);
  if (score != null) {
    const s = Math.round(score);
    qualityBadge.textContent = `${s}/100`;
    qualityBadge.style.background = s >= 80 ? 'rgba(16,185,129,0.2)' : s >= 60 ? 'rgba(245,158,11,0.2)' : 'rgba(239,68,68,0.2)';
    qualityBadge.style.color = s >= 80 ? 'var(--green)' : s >= 60 ? 'var(--amber)' : 'var(--red)';
  }
  outputCard.scrollIntoView({ behavior: 'smooth', block: 'nearest' });
}

/* ─── History ────────────────────────────────────────────────────────────── */
function saveToHistory(item) {
  item.timestamp = Date.now();
  state.taskHistory.unshift(item);
  if (state.taskHistory.length > 20) state.taskHistory.length = 20;
  localStorage.setItem('taskHistory', JSON.stringify(state.taskHistory));
  renderHistory();
}

function renderHistory() {
  if (!state.taskHistory.length) {
    taskHistory.innerHTML = '<p class="empty-state">No tasks yet</p>';
    return;
  }
  taskHistory.innerHTML = state.taskHistory.slice(0, 10).map(item => `
    <div class="history-item" data-id="${item.task_id || ''}">
      <div class="hi-task">${escHtml(item.task.slice(0, 80))}</div>
      <div class="hi-meta">
        <span class="hi-status ${item.status}">${item.status}</span>
        ${item.score != null ? `<span>${Math.round(item.score)}/100</span>` : ''}
        <span>${timeAgo(item.timestamp)}</span>
      </div>
    </div>
  `).join('');

  taskHistory.querySelectorAll('.history-item').forEach(el => {
    el.addEventListener('click', () => {
      const id = el.dataset.id;
      if (id) loadTask(id);
    });
  });
}

async function loadTask(taskId) {
  try {
    const res = await fetch(`${API_BASE}/api/tasks/${taskId}`);
    if (!res.ok) return;
    const data = await res.json();
    if (data.final_output) showOutput(data.final_output, data.quality_score);
    if (data.plan) renderPlan(data.plan);
    setStatus(data.status);
    updateMetrics({ total_tokens: data.total_tokens, plan: data.plan, iteration: null, quality_score: data.quality_score });
  } catch { /* ignore */ }
}

function timeAgo(ts) {
  const diff = (Date.now() - ts) / 1000;
  if (diff < 60) return 'just now';
  if (diff < 3600) return `${Math.floor(diff / 60)}m ago`;
  if (diff < 86400) return `${Math.floor(diff / 3600)}h ago`;
  return `${Math.floor(diff / 86400)}d ago`;
}

/* ─── Init ───────────────────────────────────────────────────────────────── */
renderHistory();
taskInput.focus();