/**
 * KOAN Demo Chat UI — shared client-side logic.
 *
 * Each demo page sets `window.DEMO_CONFIG` before loading this script:
 *   { apiBase: "/api" }
 *
 * Streams all event types: agent_start, tool_start, tool_end, token, complete, error.
 */

const API = (window.DEMO_CONFIG || {}).apiBase || '/api';

let currentThreadId = null;
let threads = [];
let isStreaming = false;

// ── DOM refs ────────────────────────────────────────────────────────

const messagesEl = document.getElementById('messages');
const inputEl = document.getElementById('input');
const sendBtn = document.getElementById('send-btn');
const threadListEl = document.getElementById('thread-list');
const newThreadBtn = document.getElementById('new-thread-btn');

// ── Thread management ───────────────────────────────────────────────

async function loadThreads() {
  try {
    const res = await fetch(`${API}/threads`);
    const data = await res.json();
    threads = data.map(t =>
      typeof t === 'string' ? { thread_id: t } : t
    );
    renderThreadList();
    if (threads.length > 0 && !currentThreadId) {
      selectThread(threads[0].thread_id);
    }
  } catch (e) {
    console.error('Failed to load threads:', e);
  }
}

function renderThreadList() {
  if (!threadListEl) return;
  threadListEl.innerHTML = '';
  for (const t of threads) {
    const el = document.createElement('div');
    const isActive = t.thread_id === currentThreadId;
    el.className = isActive
      ? 'px-3 py-2 rounded-lg text-sm cursor-pointer bg-indigo-600/20 text-indigo-300 border border-indigo-500/30'
      : 'px-3 py-2 rounded-lg text-sm cursor-pointer text-slate-400 hover:bg-white/5 hover:text-slate-200 transition';
    el.textContent = t.title || t.thread_id.slice(0, 12) + '...';
    el.onclick = () => selectThread(t.thread_id);
    threadListEl.appendChild(el);
  }
}

async function selectThread(threadId) {
  currentThreadId = threadId;
  renderThreadList();
  await loadMessages(threadId);
}

async function createThread() {
  try {
    const res = await fetch(`${API}/threads`, { method: 'POST' });
    const data = await res.json();
    threads.unshift(data);
    await selectThread(data.thread_id);
  } catch (e) {
    console.error('Failed to create thread:', e);
  }
}

async function loadMessages(threadId) {
  try {
    const res = await fetch(`${API}/threads/${threadId}/messages`);
    const messages = await res.json();
    messagesEl.innerHTML = '';
    for (const msg of messages) {
      appendMessage(msg.role, msg.content, msg.metadata || {});
    }
    scrollToBottom();
  } catch (e) {
    console.error('Failed to load messages:', e);
  }
}

// ── Message rendering ───────────────────────────────────────────────

function appendMessage(role, content, metadata = {}, live = false) {
  const isUser = role === 'user';

  // Render tool call / routing chips above the assistant bubble (skip during live streaming — already shown)
  if (!isUser && !live) {
    if (metadata.routed_to) {
      appendAgentStart(metadata.routed_to);
    }
    if (Array.isArray(metadata.tool_calls)) {
      for (const tc of metadata.tool_calls) {
        appendToolEnd(tc.tool, tc.result || '');
      }
    }
  }

  const wrapper = document.createElement('div');
  wrapper.className = `max-w-3xl mx-auto flex gap-3 ${isUser ? 'justify-end' : ''}`;

  const bubble = document.createElement('div');
  bubble.className = isUser
    ? 'bg-indigo-600/30 border border-indigo-500/30 rounded-2xl rounded-br-md px-4 py-3 max-w-[80%]'
    : 'bg-surface border border-border rounded-2xl rounded-bl-md px-4 py-3 max-w-[80%]';

  const roleEl = document.createElement('div');
  roleEl.className = 'text-xs font-semibold mb-1 ' + (isUser ? 'text-indigo-300' : 'text-accent');
  roleEl.textContent = isUser ? 'You' : 'Assistant';

  const contentEl = document.createElement('div');
  contentEl.className = 'text-sm text-slate-200 whitespace-pre-wrap break-words';
  contentEl.textContent = content;

  bubble.appendChild(roleEl);
  bubble.appendChild(contentEl);

  // Metadata tags
  if (metadata && Object.keys(metadata).length > 0) {
    const metaEl = document.createElement('div');
    metaEl.className = 'flex flex-wrap gap-1.5 mt-2';
    for (const [k, v] of Object.entries(metadata)) {
      if (k === 'run_id' || k === 'thread_id') continue;
      const tag = document.createElement('span');
      tag.className = 'text-[10px] px-2 py-0.5 rounded-full bg-white/5 text-dim ring-1 ring-white/10';
      tag.textContent = typeof v === 'object' ? `${k}: ${JSON.stringify(v)}` : `${k}: ${v}`;
      metaEl.appendChild(tag);
    }
    if (metaEl.children.length > 0) bubble.appendChild(metaEl);
  }

  // Feedback bar for assistant messages
  if (role === 'assistant' && metadata.run_id) {
    const fb = document.createElement('div');
    fb.className = 'flex gap-2 mt-2 pt-2 border-t border-border/50';
    fb.innerHTML = `
      <button class="text-xs px-2 py-1 rounded-md bg-white/5 hover:bg-emerald-500/20 text-dim hover:text-emerald-300 transition"
              onclick="submitFeedback('${metadata.run_id}', true, this)">&#128077; Helpful</button>
      <button class="text-xs px-2 py-1 rounded-md bg-white/5 hover:bg-rose-500/20 text-dim hover:text-rose-300 transition"
              onclick="submitFeedback('${metadata.run_id}', false, this)">&#128078; Not helpful</button>
    `;
    bubble.appendChild(fb);
  }

  wrapper.appendChild(bubble);
  messagesEl.appendChild(wrapper);
  return contentEl;
}

// ── Event chip (tool calls, agent routing) ──────────────────────────

function appendEventChip(html, className) {
  const wrapper = document.createElement('div');
  wrapper.className = 'max-w-3xl mx-auto';

  const chip = document.createElement('div');
  chip.className = `inline-flex items-center gap-1.5 text-[11px] px-3 py-1.5 rounded-lg ${className}`;
  chip.innerHTML = html;

  wrapper.appendChild(chip);
  messagesEl.appendChild(wrapper);
  scrollToBottom();
  return chip;
}

function appendToolStart(toolName, args) {
  const argsStr = typeof args === 'string' ? args : JSON.stringify(args);
  const truncated = argsStr.length > 80 ? argsStr.slice(0, 77) + '...' : argsStr;
  return appendEventChip(
    `<span class="inline-block w-1.5 h-1.5 rounded-full bg-amber-400 animate-pulse"></span>` +
    `<span class="font-semibold text-amber-300">${esc(toolName)}</span>` +
    `<span class="text-slate-400">${esc(truncated)}</span>`,
    'bg-amber-500/10 border border-amber-500/20'
  );
}

function appendToolEnd(toolName, result) {
  const truncated = result.length > 120 ? result.slice(0, 117) + '...' : result;
  return appendEventChip(
    `<span class="inline-block w-1.5 h-1.5 rounded-full bg-emerald-400"></span>` +
    `<span class="font-semibold text-emerald-300">${esc(toolName)}</span>` +
    `<span class="text-slate-400">${esc(truncated)}</span>`,
    'bg-emerald-500/10 border border-emerald-500/20'
  );
}

function appendAgentStart(agentName) {
  return appendEventChip(
    `<span class="inline-block w-1.5 h-1.5 rounded-full bg-violet-400 animate-pulse"></span>` +
    `<span class="font-semibold text-violet-300">Routed to</span>` +
    `<span class="text-slate-300">${esc(agentName)}</span>`,
    'bg-violet-500/10 border border-violet-500/20'
  );
}

function esc(s) {
  const d = document.createElement('div');
  d.textContent = s;
  return d.innerHTML;
}

// ── Streaming message ───────────────────────────────────────────────

function createStreamingMessage() {
  const wrapper = document.createElement('div');
  wrapper.className = 'max-w-3xl mx-auto flex gap-3';
  wrapper.id = 'streaming-msg';

  const bubble = document.createElement('div');
  bubble.className = 'bg-surface border border-border rounded-2xl rounded-bl-md px-4 py-3 max-w-[80%]';

  const roleEl = document.createElement('div');
  roleEl.className = 'text-xs font-semibold mb-1 text-accent';
  roleEl.textContent = 'Assistant';

  const contentEl = document.createElement('div');
  contentEl.className = 'text-sm text-slate-200 whitespace-pre-wrap break-words streaming-cursor';
  contentEl.id = 'streaming-content';

  bubble.appendChild(roleEl);
  bubble.appendChild(contentEl);
  wrapper.appendChild(bubble);
  messagesEl.appendChild(wrapper);
  return contentEl;
}

function scrollToBottom() {
  messagesEl.scrollTop = messagesEl.scrollHeight;
}

// ── Send message (SSE streaming) ────────────────────────────────────

async function sendMessage() {
  const text = inputEl.value.trim();
  if (!text || isStreaming) return;

  if (!currentThreadId) {
    await createThread();
  }

  // Show user message
  appendMessage('user', text);
  inputEl.value = '';
  inputEl.style.height = 'auto';
  scrollToBottom();

  // Start streaming
  isStreaming = true;
  sendBtn.disabled = true;

  let streamingContentEl = null;
  let fullText = '';
  let metadata = {};
  // Track tool call chips to update them when tool_end arrives
  const pendingTools = {};
  // Collect completed tool calls for persistence
  const toolCalls = [];
  let routedTo = null;

  try {
    const res = await fetch(`${API}/chat`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        thread_id: currentThreadId,
        message: text,
      }),
    });

    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();

      for (const line of lines) {
        if (!line.startsWith('data: ')) continue;
        const data = line.slice(6);
        if (data === '[DONE]') continue;

        let evt;
        try { evt = JSON.parse(data); } catch { continue; }

        switch (evt.type) {
          case 'agent_start':
            routedTo = evt.agent;
            appendAgentStart(evt.agent);
            break;

          case 'tool_start':
            pendingTools[evt.tool] = appendToolStart(evt.tool, evt.arguments || '');
            break;

          case 'tool_end':
            toolCalls.push({ tool: evt.tool, result: evt.result || '' });
            if (pendingTools[evt.tool]) {
              // Replace the pending chip with the completed one
              const old = pendingTools[evt.tool];
              old.querySelector('.bg-amber-400')?.classList.remove('animate-pulse');
              old.querySelector('.bg-amber-400')?.classList.replace('bg-amber-400', 'bg-emerald-400');
              // Update text
              const spans = old.querySelectorAll('span');
              if (spans[1]) spans[1].className = 'font-semibold text-emerald-300';
              const result = evt.result || '';
              const truncated = result.length > 120 ? result.slice(0, 117) + '...' : result;
              if (spans[2]) spans[2].textContent = truncated;
              old.className = old.className
                .replace('bg-amber-500/10', 'bg-emerald-500/10')
                .replace('border-amber-500/20', 'border-emerald-500/20');
              delete pendingTools[evt.tool];
            } else {
              appendToolEnd(evt.tool, evt.result || '');
            }
            break;

          case 'token':
            // Lazily create the streaming bubble on first token
            if (!streamingContentEl) {
              streamingContentEl = createStreamingMessage();
            }
            fullText += evt.token;
            streamingContentEl.textContent = fullText;
            scrollToBottom();
            break;

          case 'complete':
            metadata = evt.metadata || {};
            if (evt.run_id) metadata.run_id = evt.run_id;
            break;

          case 'error':
            if (!streamingContentEl) {
              streamingContentEl = createStreamingMessage();
            }
            fullText += `\n[Error: ${evt.error}]`;
            streamingContentEl.textContent = fullText;
            break;
        }
      }
    }
  } catch (e) {
    if (!streamingContentEl) {
      streamingContentEl = createStreamingMessage();
    }
    fullText += `\n[Connection error: ${e.message}]`;
    streamingContentEl.textContent = fullText;
  }

  // Finalize
  if (streamingContentEl) {
    streamingContentEl.classList.remove('streaming-cursor');
  }
  isStreaming = false;
  sendBtn.disabled = false;

  // Replace streaming message with final message (with metadata + feedback)
  // Inject tool_calls and routed_to so they persist and render on reload
  if (toolCalls.length > 0) metadata.tool_calls = toolCalls;
  if (routedTo) metadata.routed_to = routedTo;
  const streamMsg = document.getElementById('streaming-msg');
  if (streamMsg) streamMsg.remove();
  if (fullText) {
    // Pass _live flag so appendMessage skips chips already rendered during streaming
    appendMessage('assistant', fullText, metadata, /* live */ true);
  }
  scrollToBottom();

  // Update info panel
  updateInfoPanel(metadata);
}

// ── Feedback ────────────────────────────────────────────────────────

async function submitFeedback(runId, satisfied, btnEl) {
  try {
    await fetch(`${API}/feedback`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ run_id: runId, satisfied }),
    });
    const bar = btnEl.parentElement;
    bar.querySelectorAll('button').forEach(b => {
      b.className = 'text-xs px-2 py-1 rounded-md bg-white/5 text-dim opacity-40 cursor-not-allowed';
      b.disabled = true;
    });
    btnEl.className = satisfied
      ? 'text-xs px-2 py-1 rounded-md bg-emerald-500/20 text-emerald-300 ring-1 ring-emerald-500/30'
      : 'text-xs px-2 py-1 rounded-md bg-rose-500/20 text-rose-300 ring-1 ring-rose-500/30';
    btnEl.disabled = false;
  } catch (e) {
    console.error('Failed to submit feedback:', e);
  }
}

// ── Info panel ──────────────────────────────────────────────────────

function updateInfoPanel(metadata) {
  const panel = document.getElementById('info-content');
  if (!panel || !metadata) return;

  const kv = (key, val) =>
    `<div class="flex justify-between py-1 border-b border-border/50">` +
    `<span class="text-dim">${key}</span>` +
    `<span class="text-slate-300 font-medium">${val}</span></div>`;

  let html = '';

  if (metadata.agent_name) html += kv('Agent', metadata.agent_name);
  if (metadata.routed_to) html += kv('Routed to', metadata.routed_to);
  if (metadata.iterations !== undefined) html += kv('Iterations', metadata.iterations);
  if (metadata.tool_calls_made !== undefined) html += kv('Tool calls', metadata.tool_calls_made);
  if (metadata.run_id) html += kv('Run ID', metadata.run_id.slice(0, 8) + '...');

  if (metadata.graph_context) {
    const gc = metadata.graph_context;
    html += kv('Precedents', gc.precedents || 0);
    html += kv('Entities', gc.entities || 0);
    html += kv('Policies', gc.policies || 0);
    html += kv('Has feedback', gc.has_feedback || false);
  }

  if (metadata.evaluation) {
    const ev = metadata.evaluation;
    html += kv('Review status', ev.review_status);
    html += kv('Confidence', ev.calibrated_confidence);
  }

  if (metadata.calibration) {
    const cal = metadata.calibration;
    html += kv('Cal. factor', cal.factor?.toFixed(3) || 'N/A');
    html += kv('Samples', cal.sample_count || 0);
  }

  if (html) panel.innerHTML = html;
}

// ── Input handling ──────────────────────────────────────────────────

inputEl.addEventListener('keydown', (e) => {
  if (e.key === 'Enter' && !e.shiftKey) {
    e.preventDefault();
    sendMessage();
  }
});

inputEl.addEventListener('input', () => {
  inputEl.style.height = 'auto';
  inputEl.style.height = Math.min(inputEl.scrollHeight, 120) + 'px';
});

sendBtn.addEventListener('click', sendMessage);
if (newThreadBtn) newThreadBtn.addEventListener('click', createThread);

// ── Init ────────────────────────────────────────────────────────────
loadThreads();
