const msgs = document.getElementById('messages');
const inputEl = document.getElementById('input');
const statusEl = document.getElementById('status');
const traceEl = document.getElementById('trace');
let ws, currentEl;
let authToken = localStorage.getItem('cog_token');
let authConfig = null;

// --- OIDC Auth ---

async function initAuth() {
  try {
    const resp = await fetch('/auth/config');
    authConfig = await resp.json();
  } catch { authConfig = { enabled: false }; }

  if (!authConfig.enabled) { connect(); return; }

  // Handle OIDC callback
  if (location.pathname === '/callback') {
    const params = new URLSearchParams(location.search);
    const code = params.get('code');
    const verifier = sessionStorage.getItem('pkce_verifier');
    if (code && verifier) {
      const tokenResp = await fetch(authConfig.issuer + '/oauth/v2/token', {
        method: 'POST',
        headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
        body: new URLSearchParams({
          grant_type: 'authorization_code',
          client_id: authConfig.clientId,
          code,
          redirect_uri: location.origin + '/callback',
          code_verifier: verifier,
        }),
      });
      const tokens = await tokenResp.json();
      if (tokens.access_token) {
        // Store access token for userinfo, id_token for JWT validation
        localStorage.setItem('cog_access_token', tokens.access_token);
        authToken = tokens.id_token || tokens.access_token;
        localStorage.setItem('cog_token', authToken);
        sessionStorage.removeItem('pkce_verifier');
      }
    }
    history.replaceState(null, '', '/');
  }

  if (authToken) {
    connect();
  } else {
    showLogin();
  }
}

function showLogin() {
  statusEl.textContent = 'not authenticated';
  statusEl.style.color = '#f59e0b';
  const btn = document.createElement('button');
  btn.textContent = 'Log in with loop42';
  btn.className = 'login-btn';
  btn.onclick = startLogin;
  document.getElementById('input-bar').replaceChildren(btn);
}

async function startLogin() {
  // PKCE: generate code_verifier + code_challenge
  const verifier = randomString(64);
  sessionStorage.setItem('pkce_verifier', verifier);
  const encoder = new TextEncoder();
  const digest = await crypto.subtle.digest('SHA-256', encoder.encode(verifier));
  const challenge = btoa(String.fromCharCode(...new Uint8Array(digest)))
    .replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '');

  const params = new URLSearchParams({
    response_type: 'code',
    client_id: authConfig.clientId,
    redirect_uri: location.origin + '/callback',
    scope: 'openid profile email',
    code_challenge: challenge,
    code_challenge_method: 'S256',
  });
  location.href = authConfig.issuer + '/oauth/v2/authorize?' + params;
}

function randomString(len) {
  const arr = new Uint8Array(len);
  crypto.getRandomValues(arr);
  return btoa(String.fromCharCode(...arr)).replace(/[^a-zA-Z0-9]/g, '').slice(0, len);
}

// --- WebSocket ---

function connect() {
  const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
  let wsUrl = proto + '//' + location.host + '/ws';
  if (authToken) {
    const accessToken = localStorage.getItem('cog_access_token') || '';
    wsUrl += '?token=' + encodeURIComponent(authToken) + '&access_token=' + encodeURIComponent(accessToken);
  }
  ws = new WebSocket(wsUrl);

  ws.onopen = () => {
    statusEl.textContent = 'connected';
    statusEl.style.color = '#22c55e';
    addTrace('runtime', 'connected', 'ws open');
  };

  ws.onclose = () => {
    statusEl.textContent = 'disconnected';
    statusEl.style.color = '#666';
    addTrace('runtime', 'disconnected', 'ws closed');
    setTimeout(connect, 2000);
  };

  ws.onmessage = (e) => {
    const data = JSON.parse(e.data);

    if (data.type === 'hud') {
      handleHud(data);

    } else if (data.type === 'delta') {
      if (!currentEl) {
        currentEl = addMsg('assistant', '');
        currentEl.classList.add('streaming');
      }
      currentEl.textContent += data.content;
      scroll(msgs);

    } else if (data.type === 'done') {
      if (currentEl) currentEl.classList.remove('streaming');
      currentEl = null;
    }
  };
}

function handleHud(data) {
  const node = data.node || 'unknown';
  const event = data.event || '';

  if (event === 'context') {
    // Update node meter
    if (data.tokens !== undefined) {
      updateMeter(node, data.tokens, data.max_tokens, data.fill_pct);
    }
    // Expandable: show message count + token info
    const count = (data.messages || []).length;
    const tokenInfo = data.tokens ? ` [${data.tokens}/${data.max_tokens}t ${data.fill_pct}%]` : '';
    const summary = count + ' msgs' + tokenInfo + ': ' + (data.messages || []).map(m =>
      m.role[0].toUpperCase() + ':' + truncate(m.content, 30)
    ).join(' | ');
    const detail = (data.messages || []).map((m, i) =>
      i + ' [' + m.role + '] ' + m.content
    ).join('\n');
    addTrace(node, 'context', summary, 'context', detail);

  } else if (event === 'perceived') {
    addTrace(node, 'perceived', data.instruction, 'instruction');

  } else if (event === 'decided') {
    addTrace(node, 'decided', data.instruction, 'instruction');

  } else if (event === 'updated' && data.state) {
    const pairs = Object.entries(data.state).map(([k, v]) => {
      const val = Array.isArray(v) ? v.join(', ') : (v || 'null');
      return k + '=' + truncate(val, 25);
    }).join('  ');
    const detail = JSON.stringify(data.state, null, 2);
    addTrace(node, 'state', pairs, 'state', detail);

  } else if (event === 'error') {
    addTrace(node, 'error', data.detail || '', 'error');

  } else if (event === 'thinking') {
    addTrace(node, 'thinking', data.detail || '');

  } else if (event === 'streaming') {
    addTrace(node, 'streaming', '');

  } else if (event === 'done') {
    addTrace(node, 'done', '');

  } else if (event === 'tick') {
    // Update sensor meter with tick count
    const meter = document.getElementById('meter-sensor');
    if (meter) {
      const text = meter.querySelector('.nm-text');
      const deltas = Object.entries(data.deltas || {}).map(([k,v]) => k + '=' + v).join(' ');
      text.textContent = 'tick #' + (data.tick || 0) + (deltas ? ' | ' + deltas : '');
    }
    if (data.deltas && Object.keys(data.deltas).length) {
      const deltas = Object.entries(data.deltas).map(([k,v]) => k + '=' + truncate(String(v), 30)).join(' ');
      addTrace(node, 'tick #' + data.tick, deltas);
    }

  } else if (event === 'started' || event === 'stopped') {
    const meter = document.getElementById('meter-sensor');
    if (meter) meter.querySelector('.nm-text').textContent = event;
    addTrace(node, event, '');

  } else {
    // Generic fallback
    const detail = JSON.stringify(data, null, 2);
    addTrace(node, event, '', '', detail);
  }
}

function addTrace(node, event, text, cls, detail) {
  const line = document.createElement('div');
  line.className = 'trace-line' + (detail ? ' expandable' : '');

  const ts = new Date().toLocaleTimeString('de-DE', { hour12: false, hour: '2-digit', minute: '2-digit', second: '2-digit', fractionalSecondDigits: 1 });

  line.innerHTML =
    '<span class="trace-ts">' + ts + '</span>' +
    '<span class="trace-node ' + esc(node) + '">' + esc(node) + '</span>' +
    '<span class="trace-event">' + esc(event) + '</span>' +
    '<span class="trace-data' + (cls ? ' ' + cls : '') + '">' + esc(text) + '</span>';

  traceEl.appendChild(line);

  if (detail) {
    const detailEl = document.createElement('div');
    detailEl.className = 'trace-detail';
    detailEl.textContent = detail;
    traceEl.appendChild(detailEl);
    line.addEventListener('click', () => detailEl.classList.toggle('open'));
  }

  scroll(traceEl);
}

function updateMeter(node, tokens, maxTokens, fillPct) {
  const meter = document.getElementById('meter-' + node);
  if (!meter) return;
  const fill = meter.querySelector('.nm-fill');
  const text = meter.querySelector('.nm-text');
  fill.style.width = fillPct + '%';
  fill.style.backgroundColor = fillPct > 80 ? '#ef4444' : fillPct > 50 ? '#f59e0b' : '#22c55e';
  text.textContent = tokens + ' / ' + maxTokens + 't (' + fillPct + '%)';
}

function scroll(el) { el.scrollTop = el.scrollHeight; }
function esc(s) { const d = document.createElement('span'); d.textContent = s; return d.innerHTML; }
function truncate(s, n) { return s.length > n ? s.slice(0, n) + '\u2026' : s; }

function addMsg(role, text) {
  const el = document.createElement('div');
  el.className = 'msg ' + role;
  el.textContent = text;
  msgs.appendChild(el);
  scroll(msgs);
  return el;
}

function send() {
  const text = inputEl.value.trim();
  if (!text || !ws || ws.readyState !== 1) return;
  addMsg('user', text);
  addTrace('runtime', 'user_msg', truncate(text, 60));
  ws.send(JSON.stringify({ text }));
  inputEl.value = '';
}

inputEl.addEventListener('keydown', (e) => { if (e.key === 'Enter') send(); });
initAuth();