public/js/chat.js

// chat.js — Chat rendering, streaming, versioning, editing
import { send, on, off } from './ws.js';
import { currentSessionId } from './sessions.js';
import {
  renderMarkdown, attachCodeCopyListeners, attachSvgPanelListeners,
  escHtml, showNotification, autoResize,
} from './ui.js';

let activeSessionId = null;
let isStreaming      = false;
let streamingBubble  = null;
let streamingText    = '';
let autoScroll       = true;

export function setActiveSession(id) { activeSessionId = id; }
export function getIsStreaming()     { return isStreaming; }

// ── WebSocket events ──────────────────────────────────────────────────────

on('sessions:data',         (msg) => { if (msg.session.id === activeSessionId) renderSession(msg.session); });
on('chat:start',            (msg) => { if (msg.sessionId === activeSessionId) onChatStart(); });
on('chat:token',            (msg) => { if (msg.sessionId === activeSessionId) onToken(msg.token); });
on('chat:done',             (msg) => { if (msg.sessionId === activeSessionId) onChatDone(msg); });
on('chat:aborted',          (msg) => { if (msg.sessionId === activeSessionId) onChatAborted(msg); });
on('chat:error',            (msg) => { if (msg.sessionId === activeSessionId) onChatError(msg.error); });
on('chat:asset',            (msg) => { if (msg.sessionId === activeSessionId) appendAsset(msg.asset); });
on('chat:toolCall',         (msg) => { if (msg.sessionId === activeSessionId) handleLiveToolCall(msg.call); });
on('chat:messageEdited',    (msg) => { if (msg.sessionId === activeSessionId) renderHistory(msg.history); });
on('chat:versionSelected',  (msg) => { if (msg.sessionId === activeSessionId) renderHistory(msg.history); });

// ── Views ─────────────────────────────────────────────────────────────────

export function renderSession(session) {
  if (!session || !session.history?.length) { showWelcome(); return; }
  showChat();
  renderHistory(session.history);
}

function showWelcome() {
  document.getElementById('welcome-view')?.classList.remove('hidden');
  document.getElementById('chat-view')?.classList.add('hidden');
  document.getElementById('bottom-input-bar')?.classList.add('hidden');
}

function showChat() {
  document.getElementById('welcome-view')?.classList.add('hidden');
  document.getElementById('chat-view')?.classList.remove('hidden');
  document.getElementById('bottom-input-bar')?.classList.remove('hidden');
}

// ── Full history render ───────────────────────────────────────────────────

export function renderHistory(history) {
  showChat();
  const box = document.getElementById('chat-messages');
  if (!box) return;
  box.innerHTML = '';

  for (let i = 0; i < history.length; i++) {
    const msg = history[i];
    if      (msg.role === 'user')      appendUserMsg(box, msg, i);
    else if (msg.role === 'assistant') appendAssistantMsg(box, msg, i);
    else if (msg.role === 'image')     appendMediaMsg(box, 'image', msg.content);
    else if (msg.role === 'video')     appendMediaMsg(box, 'video', msg.content);
    else if (msg.role === 'audio')     appendMediaMsg(box, 'audio', msg.content);
  }

  // Render math if KaTeX is available
  if (typeof renderMathInElement !== 'undefined') {
    try {
      renderMathInElement(box, {
        delimiters: [
          { left: '$$', right: '$$', display: true },
          { left: '$',  right: '$',  display: false },
          { left: '\\(', right: '\\)', display: false },
          { left: '\\[', right: '\\]', display: true },
        ],
        throwOnError: false,
      });
    } catch {}
  }

  if (autoScroll) box.scrollTop = box.scrollHeight;
}

// ── Message renderers ─────────────────────────────────────────────────────

function appendUserMsg(box, msg, index) {
  const wrap = makeWrap(index);
  const bubble = document.createElement('div');
  bubble.className = 'msg-user';

  const text = msgText(msg.content);
  const imgs = msgImages(msg.content);

  bubble.innerHTML = renderMarkdown(text);
  attachCodeCopyListeners(bubble);
  attachSvgPanelListeners(bubble);

  imgs.forEach(src => {
    const img = document.createElement('img');
    img.src = src; img.alt = 'Attached image';
    img.style.cssText = 'max-width:100%;max-height:260px;border-radius:8px;margin-top:8px;display:block;cursor:pointer;';
    img.addEventListener('click', () => openImageModal(src));
    bubble.appendChild(img);
  });

  wrap.appendChild(bubble);
  if (msg.versions?.length > 1) wrap.appendChild(buildVersionNav(msg, index));
  wrap.appendChild(buildActions([
    { icon: '📋', title: 'Copy', fn: () => copyText(text) },
    { icon: '✏️', title: 'Edit', fn: () => startUserEdit(wrap, index, msg, text) },
  ]));
  box.appendChild(wrap);
}

function appendAssistantMsg(box, msg, index) {
  const wrap = makeWrap(index);
  const bubble = document.createElement('div');
  bubble.className = 'msg-assistant';
  bubble.innerHTML = renderMarkdown(msg.content || '');
  attachCodeCopyListeners(bubble);
  attachSvgPanelListeners(bubble);

  // Tool call chips
  if (msg.toolCalls?.length) {
    const chipRow = document.createElement('div');
    chipRow.style.cssText = 'display:flex;flex-wrap:wrap;gap:4px;margin-top:6px;';
    msg.toolCalls.forEach(c => chipRow.appendChild(buildToolChip(c)));
    bubble.appendChild(chipRow);
  }

  wrap.appendChild(bubble);
  if (msg.versions?.length > 1) wrap.appendChild(buildVersionNav(msg, index));
  wrap.appendChild(buildActions([
    { icon: '📋', title: 'Copy', fn: () => copyText(msg.content || '') },
    { icon: '✏️', title: 'Edit', fn: () => startAssistantEdit(wrap, index, msg) },
  ]));
  box.appendChild(wrap);
}

function appendMediaMsg(box, type, content) {
  const wrap = document.createElement('div');
  wrap.className = 'msg-media';

  if (type === 'image') {
    const img = document.createElement('img');
    img.src = content; img.alt = 'Generated image';
    img.addEventListener('click', () => openImageModal(content));
    wrap.appendChild(img);
    wrap.appendChild(dlBtn(() => dlMedia(content, 'image.png')));
  } else if (type === 'video') {
    const v = document.createElement('video');
    v.src = content; v.controls = true; v.preload = 'metadata';
    wrap.appendChild(v);
    wrap.appendChild(dlBtn(() => dlMedia(content, 'video.mp4')));
  } else if (type === 'audio') {
    const a = document.createElement('audio');
    a.src = content; a.controls = true; a.preload = 'metadata'; a.style.width = '100%';
    wrap.appendChild(a);
    const db = dlBtn(() => dlMedia(content, 'audio.mp3'));
    db.style.cssText += ';position:static;margin-top:6px;opacity:1;';
    wrap.appendChild(db);
  }
  box.appendChild(wrap);
}

// ── Builders ──────────────────────────────────────────────────────────────

function makeWrap(index) {
  const w = document.createElement('div');
  w.className = 'msg-group'; w.dataset.index = index;
  return w;
}

function buildActions(items) {
  const div = document.createElement('div');
  div.className = 'msg-actions';
  items.forEach(({ icon, title, fn }) => {
    const btn = document.createElement('button');
    btn.className = 'msg-action-btn'; btn.title = title;
    btn.innerHTML = icon; btn.style.fontSize = '13px';
    btn.addEventListener('click', e => { e.stopPropagation(); fn(); });
    div.appendChild(btn);
  });
  return div;
}

function buildVersionNav(msg, index) {
  const nav     = document.createElement('div');
  nav.className = 'msg-version-nav';
  const total   = msg.versions.length;
  const cur     = (msg.currentVersionIdx ?? 0) + 1;

  const prev = document.createElement('button');
  prev.textContent = '‹'; prev.title = 'Previous version'; prev.disabled = cur <= 1;
  prev.addEventListener('click', () => send({ type: 'chat:selectVersion',
    sessionId: activeSessionId, messageIndex: index, versionIdx: (msg.currentVersionIdx ?? 0) - 1 }));

  const lbl = document.createElement('span');
  lbl.textContent = `${cur} / ${total}`;
  lbl.style.cssText = 'min-width:36px;text-align:center;';

  const next = document.createElement('button');
  next.textContent = '›'; next.title = 'Next version'; next.disabled = cur >= total;
  next.addEventListener('click', () => send({ type: 'chat:selectVersion',
    sessionId: activeSessionId, messageIndex: index, versionIdx: (msg.currentVersionIdx ?? 0) + 1 }));

  nav.appendChild(prev); nav.appendChild(lbl); nav.appendChild(next);
  return nav;
}

function buildToolChip(call) {
  const names = { ollama_search: 'Web Search', read_web_page: 'Read Page',
    generate_image: 'Image Gen', generate_video: 'Video Gen', generate_audio: 'Audio Gen' };
  const icons = { ollama_search: '🔍', read_web_page: '📄',
    generate_image: '🖼️', generate_video: '🎬', generate_audio: '🎵' };
  const chip = document.createElement('button');
  chip.className = 'msg-tool-call';
  chip.innerHTML = `<span>${icons[call.name] || '🔧'}</span><span>${escHtml(names[call.name] || call.name)}</span>`;
  chip.addEventListener('click', () => import('./modals.js').then(m => m.showToolCallModal(call)));
  return chip;
}

function dlBtn(fn) {
  const btn = document.createElement('button');
  btn.className = 'media-download-btn'; btn.textContent = 'Download';
  btn.addEventListener('click', fn);
  return btn;
}

// ── Inline editing ────────────────────────────────────────────────────────

function startUserEdit(wrap, index, msg, originalText) {
  const bubble = wrap.querySelector('.msg-user'); if (!bubble) return;
  bubble.innerHTML = '';

  const ta = makeEditTextarea(originalText);
  bubble.appendChild(ta);

  const btns = document.createElement('div');
  btns.style.cssText = 'display:flex;gap:6px;margin-top:8px;justify-content:flex-end;';

  const cancelBtn = makeBtn('Cancel', 'btn-ghost');
  cancelBtn.addEventListener('click', () => {
    bubble.innerHTML = renderMarkdown(originalText);
    attachCodeCopyListeners(bubble); btns.remove();
  });

  const sendBtn = makeBtn('Send', 'btn-primary');
  sendBtn.addEventListener('click', () => {
    const newContent = ta.value.trim(); if (!newContent) return;
    // Store original session tail in new version, then server truncates history to this index
    send({ type: 'chat:editMessage', sessionId: activeSessionId,
      messageIndex: index, newContent, role: 'user' });

    // After edit confirmed, re-send the new content as a new message
    const handler = (editMsg) => {
      if (editMsg.sessionId !== activeSessionId || editMsg.messageIndex !== index) return;
      off('chat:messageEdited', handler);
      // Re-submit from this message
      import('./app.js').catch(() => {}).finally(() => {
        send({ type: 'chat:send', sessionId: activeSessionId,
          content: newContent, tools: getActiveTools() });
      });
    };
    on('chat:messageEdited', handler);
  });

  btns.appendChild(cancelBtn); btns.appendChild(sendBtn);
  bubble.appendChild(btns);
  ta.focus();
}

function startAssistantEdit(wrap, index, msg) {
  const bubble = wrap.querySelector('.msg-assistant'); if (!bubble) return;
  const originalContent = msg.content || '';
  bubble.classList.add('editing');
  bubble.innerHTML = '';

  const ta = makeEditTextarea(originalContent, 20);
  bubble.appendChild(ta);

  const btns = document.createElement('div');
  btns.className = 'edit-actions';

  const cancelBtn = makeBtn('Cancel', 'btn-ghost');
  cancelBtn.addEventListener('click', () => {
    bubble.classList.remove('editing');
    bubble.innerHTML = renderMarkdown(originalContent);
    attachCodeCopyListeners(bubble);
    btns.remove();
  });

  const saveBtn = makeBtn('Save', 'btn-primary');
  saveBtn.addEventListener('click', () => {
    const newContent = ta.value.trim(); if (!newContent) return;
    send({ type: 'chat:editMessage', sessionId: activeSessionId,
      messageIndex: index, newContent, role: 'assistant' });
  });

  btns.appendChild(cancelBtn); btns.appendChild(saveBtn);
  bubble.appendChild(btns);
  ta.focus();
}

function makeEditTextarea(value, maxLines = 8) {
  const ta = document.createElement('textarea');
  ta.value = value;
  ta.style.cssText = 'width:100%;background:transparent;border:none;outline:none;' +
    'color:var(--text);font:inherit;font-size:14px;resize:none;line-height:1.6;';
  autoResize(ta, maxLines);
  ta.addEventListener('input', () => autoResize(ta, maxLines));
  return ta;
}

function makeBtn(label, cls) {
  const btn = document.createElement('button');
  btn.textContent = label; btn.className = cls;
  btn.style.cssText += ';font-size:12px;padding:5px 12px;';
  return btn;
}

// ── Streaming ─────────────────────────────────────────────────────────────

function onChatStart() {
  isStreaming = true; streamingText = '';
  showChat();
  const box = document.getElementById('chat-messages'); if (!box) return;
  streamingBubble = document.createElement('div');
  streamingBubble.className = 'msg-assistant msg-generating';

  const thinking = document.createElement('div'); thinking.className = 'msg-thinking';
  for (let i = 0; i < 3; i++) {
    const d = document.createElement('div'); d.className = 'thinking-dot'; thinking.appendChild(d);
  }
  streamingBubble.appendChild(thinking);
  box.appendChild(streamingBubble);
  if (autoScroll) box.scrollTop = box.scrollHeight;
  updateSendBtn(true);
}

function onToken(token) {
  if (!streamingBubble) return;
  streamingText += token;
  streamingBubble.querySelector('.msg-thinking')?.remove();
  streamingBubble.innerHTML = renderMarkdown(processDisplay(streamingText));
  attachCodeCopyListeners(streamingBubble);
  if (autoScroll) {
    const box = document.getElementById('chat-messages');
    if (box) box.scrollTop = box.scrollHeight;
  }
}

function onChatDone(msg) {
  isStreaming = false;
  streamingBubble?.classList.remove('msg-generating');
  streamingBubble = null; streamingText = '';
  updateSendBtn(false);
  if (msg.history) renderHistory(msg.history);
}

function onChatAborted(msg) {
  isStreaming = false;
  if (streamingBubble) {
    streamingBubble.classList.remove('msg-generating');
    const note = document.createElement('div');
    note.style.cssText = 'font-size:12px;color:var(--text-muted);margin-top:6px;';
    note.textContent = '⚠ Interrupted';
    streamingBubble.appendChild(note);
    streamingBubble = null;
  }
  updateSendBtn(false);
  if (msg.history) renderHistory(msg.history);
}

function onChatError(err) {
  isStreaming = false;
  if (streamingBubble) {
    streamingBubble.classList.remove('msg-generating');
    streamingBubble.querySelector('.msg-thinking')?.remove();
    const note = document.createElement('div');
    note.style.cssText = 'color:#f87171;font-size:13px;margin-top:6px;';
    note.textContent = `⚠ Error: ${err}`;
    streamingBubble.appendChild(note);
    streamingBubble = null;
  }
  updateSendBtn(false);
}

function handleLiveToolCall(call) {
  if (!streamingBubble) return;
  const names = { ollama_search: 'Searching web…', read_web_page: 'Reading page…',
    generate_image: 'Generating image…', generate_video: 'Generating video…', generate_audio: 'Generating audio…' };

  if (call.state === 'pending') {
    streamingBubble.querySelector('.msg-thinking')?.remove();
    // Don't add duplicates
    if (!streamingBubble.querySelector(`[data-tcid="${call.id}"]`)) {
      const badge = document.createElement('div'); badge.className = 'msg-tool-call';
      badge.style.pointerEvents = 'none'; badge.setAttribute('data-tcid', call.id);
      badge.innerHTML = `<span>🔧</span><span>${names[call.name] || call.name}</span>`;
      streamingBubble.appendChild(badge);
    }
  } else if (call.state === 'resolved' || call.state === 'canceled') {
    streamingBubble.querySelector(`[data-tcid="${call.id}"]`)?.remove();
  }
}

function appendAsset(asset) {
  const box = document.getElementById('chat-messages'); if (!box) return;
  appendMediaMsg(box, asset.role, asset.content);
  if (autoScroll) box.scrollTop = box.scrollHeight;
}

// ── Submit ────────────────────────────────────────────────────────────────

export function submitMessage(text, attachments = []) {
  if (!text.trim() && attachments.length === 0) return;
  if (isStreaming) { send({ type: 'chat:stop' }); return; }
  if (!activeSessionId) return;

  const images   = attachments.filter(a => a.type === 'image');
  const textFiles= attachments.filter(a => a.type === 'text');

  let fullText = text;
  if (textFiles.length > 0) {
    fullText += '\n\n<details><summary>Attached Files</summary>\n';
    for (const f of textFiles)
      fullText += `\n<details><summary>${f.name}</summary>\n\n\`\`\`\n${f.content}\n\`\`\`\n\n</details>\n`;
    fullText += '</details>';
  }

  let content;
  if (images.length > 0) {
    content = [
      { type: 'text', text: fullText },
      ...images.map(img => ({ type: 'image_url', image_url: { url: `data:${img.mimeType};base64,${img.base64}` } })),
    ];
  } else {
    content = fullText;
  }

  // Append optimistic user bubble
  const box = document.getElementById('chat-messages');
  if (box) {
    const wrap = makeWrap(-1);
    const bubble = document.createElement('div'); bubble.className = 'msg-user';
    bubble.innerHTML = renderMarkdown(fullText);
    images.forEach(img => {
      const el = document.createElement('img');
      el.src = `data:${img.mimeType};base64,${img.base64}`;
      el.style.cssText = 'max-width:100%;max-height:200px;border-radius:8px;margin-top:6px;display:block;';
      bubble.appendChild(el);
    });
    wrap.appendChild(bubble);
    box.appendChild(wrap);
    if (autoScroll) box.scrollTop = box.scrollHeight;
  }

  send({ type: 'chat:send', sessionId: activeSessionId, content, tools: getActiveTools(), clientId: localStorage.getItem('ipai_client_id') || '' });
}

// ── Utils ─────────────────────────────────────────────────────────────────

function getActiveTools() {
  const tools = {};
  document.querySelectorAll('[data-tool]').forEach(btn => {
    if (btn.dataset.tool) tools[btn.dataset.tool] = btn.classList.contains('active');
  });
  return tools;
}

function msgText(content) {
  if (typeof content === 'string') return content;
  return content.filter(p => p.type === 'text').map(p => p.text).join('\n');
}
function msgImages(content) {
  if (typeof content === 'string') return [];
  return content.filter(p => p.type === 'image_url').map(p => p.image_url.url);
}

function processDisplay(text) {
  // Replace open ```svg blocks with placeholder so markdown doesn't choke
  let result = '', i = 0;
  while (i < text.length) {
    const start = text.indexOf('```svg', i);
    if (start === -1) { result += text.slice(i); break; }
    result += text.slice(i, start) + '[SVG Image]';
    const end = text.indexOf('```', start + 6);
    if (end === -1) break;
    i = end + 3;
  }
  return result;
}

function updateSendBtn(streaming) {
  document.querySelectorAll('#bottom-send-btn, #center-send-btn').forEach(btn => {
    btn.innerHTML = streaming
      ? '<svg width="16" height="16" viewBox="0 0 24 24" fill="currentColor"><rect x="5" y="5" width="14" height="14" rx="2"/></svg>'
      : '<svg width="16" height="16" viewBox="0 0 24 24" fill="currentColor"><path d="M2 21L23 12 2 3v7l15 2-15 2v7z"/></svg>';
    btn.classList.toggle('stop', streaming);
  });
}

function copyText(text) {
  navigator.clipboard.writeText(text).then(
    () => showNotification({ type: 'success', message: 'Copied', duration: 1500 }),
    () => showNotification({ type: 'error',   message: 'Copy failed', duration: 1500 })
  );
}

function dlMedia(dataUrl, filename) {
  const a = document.createElement('a');
  a.href = dataUrl; a.download = filename;
  document.body.appendChild(a); a.click(); document.body.removeChild(a);
}

function openImageModal(src) {
  import('./modals.js').then(m => m.openImageModal(src));
}

// Scroll tracking
document.getElementById('chat-view')?.addEventListener('scroll', e => {
  const el = e.target;
  autoScroll = el.scrollHeight - el.scrollTop - el.clientHeight < 60;
}, true);