#!/usr/bin/env node
// Measure submit (Enter) latency in the composer.
//
// For each round:
//   1. Focus composer, type N chars of stub text
//   2. Mark a timestamp, fire Enter via Input.dispatchKeyEvent
//   3. Observe: time until the composer becomes empty (submit accepted),
//      time until the user message renders in the thread viewport,
//      time until the optional "running…" indicator appears,
//      time until the next frame is painted after the message renders.
//
// Pre-condition: a session is loaded (load via click-session.mjs first).
// Note: this DOES talk to the real gateway/agent, so each round triggers
//       a real prompt submission. Don't run this on a live conversation
//       you care about — use a throwaway session.

import { writeFileSync } from 'node:fs'

const args = Object.fromEntries(
  process.argv.slice(2).flatMap(s => {
    const m = s.match(/^--([^=]+)(?:=(.*))?$/)
    return m ? [[m[1], m[2] ?? true]] : []
  })
)
const PORT = Number(args.port ?? 9222)
const ROUNDS = Number(args.rounds ?? 3)

async function pickRenderer() {
  const list = await (await fetch(`http://127.0.0.1:${PORT}/json/list`)).json()
  return list.find(t => t.type === 'page' && t.url.startsWith('http'))
}

function connect(url) {
  return new Promise((resolve, reject) => {
    const ws = new WebSocket(url)
    let id = 0
    const pending = new Map()
    ws.addEventListener('open', () =>
      resolve({
        send(method, params = {}) {
          const myId = ++id
          ws.send(JSON.stringify({ id: myId, method, params }))
          return new Promise((res, rej) => pending.set(myId, { res, rej }))
        },
        close: () => ws.close()
      })
    )
    ws.addEventListener('error', reject)
    ws.addEventListener('message', ev => {
      const m = JSON.parse(typeof ev.data === 'string' ? ev.data : ev.data.toString('utf8'))
      if (m.id != null) {
        const p = pending.get(m.id)
        if (!p) return
        pending.delete(m.id)
        m.error ? p.rej(new Error(m.error.message)) : p.res(m.result)
      }
    })
  })
}

async function evalP(cdp, expr) {
  const r = await cdp.send('Runtime.evaluate', { expression: expr, returnByValue: true, awaitPromise: true })
  if (r.exceptionDetails) throw new Error(r.exceptionDetails.text)
  return r.result.value
}

async function focusAndType(cdp, text) {
  await evalP(cdp, `
    (() => {
      const el = document.querySelector('[data-slot="composer-rich-input"]')
      if (!el) return
      el.focus()
      const range = document.createRange()
      range.selectNodeContents(el)
      range.collapse(false)
      const sel = window.getSelection()
      sel.removeAllRanges()
      sel.addRange(range)
    })()
  `)
  for (const c of text) {
    await cdp.send('Input.dispatchKeyEvent', { type: 'char', text: c, unmodifiedText: c })
    await new Promise(r => setTimeout(r, 8))
  }
}

async function submitAndMeasure(cdp, timeoutMs = 5000) {
  // Install observers, await milestones. Returns a promise that we'll resolve
  // after the page-side Promise resolves OR a CDP-side timeout (belt + braces).
  const setup = await evalP(
    cdp,
    `
    (() => {
      const composer = document.querySelector('[data-slot="composer-rich-input"]')
      const threadRoot = document.querySelector('[data-slot="aui_thread-content"]') ||
                         document.querySelector('[data-slot="aui_thread-viewport"]')
      const startMessageCount = threadRoot ? threadRoot.querySelectorAll('[data-slot="aui_turn-pair"], [data-slot="aui_message"]').length : 0

      window.__submitMilestones = { start: performance.now(), startMessageCount }
      window.__submitDone = false
      window.__submitResolve = null

      const composerObs = new MutationObserver(() => {
        const m = window.__submitMilestones
        if (!m) return
        if (!m.composerClearedMs && composer && composer.innerText.length === 0) {
          m.composerClearedMs = performance.now() - m.start
        }
      })
      composer && composerObs.observe(composer, { childList: true, subtree: true, characterData: true })

      const threadObs = threadRoot ? new MutationObserver(() => {
        const m = window.__submitMilestones
        if (!m || m.userMessageRenderedMs) return
        const c = threadRoot.querySelectorAll('[data-slot="aui_turn-pair"], [data-slot="aui_message"]').length
        if (c > m.startMessageCount) {
          m.userMessageRenderedMs = performance.now() - m.start
          requestAnimationFrame(() => {
            m.userMessagePaintMs = performance.now() - m.start
          })
        }
      }) : null
      threadObs && threadObs.observe(threadRoot, { childList: true, subtree: true })

      window.__submitObservers = { composerObs, threadObs }
      return true
    })()
    `
  )
  if (!setup) throw new Error('observer setup failed')

  // Send Enter via real keystroke channel (rawKeyDown + char + keyUp).
  // React synthetic onKeyDown receives this exactly like a hardware Enter.
  await cdp.send('Input.dispatchKeyEvent', {
    type: 'rawKeyDown',
    windowsVirtualKeyCode: 13,
    nativeVirtualKeyCode: 13,
    key: 'Enter',
    code: 'Enter',
    text: '\r',
    unmodifiedText: '\r'
  })
  await cdp.send('Input.dispatchKeyEvent', {
    type: 'keyUp',
    windowsVirtualKeyCode: 13,
    nativeVirtualKeyCode: 13,
    key: 'Enter',
    code: 'Enter'
  })

  // Poll for the milestones from outside; cap at timeoutMs.
  const deadline = Date.now() + timeoutMs
  while (Date.now() < deadline) {
    const m = await evalP(cdp, 'JSON.stringify(window.__submitMilestones)')
    const parsed = JSON.parse(m || '{}')
    if (parsed.userMessagePaintMs != null) {
      parsed.reason = 'paint'
      await evalP(cdp, `(() => {
        window.__submitObservers?.composerObs.disconnect()
        window.__submitObservers?.threadObs?.disconnect()
      })()`)
      return parsed
    }
    await new Promise(r => setTimeout(r, 50))
  }
  // Timed out
  const m = await evalP(cdp, 'JSON.stringify(window.__submitMilestones)')
  const parsed = JSON.parse(m || '{}')
  parsed.reason = 'timeout-overall'
  await evalP(cdp, `(() => {
    window.__submitObservers?.composerObs.disconnect()
    window.__submitObservers?.threadObs?.disconnect()
  })()`)
  return parsed
}

async function tryCancel(cdp) {
  // Find a Stop / Cancel button and click it. After submit, the composer
  // turns into "cancel" mode; clicking it interrupts the agent turn so we
  // don't burn tokens on these probes.
  await evalP(
    cdp,
    `
    (() => {
      // Common selectors: aria-label="Stop response", data-slot="composer-cancel",
      // role=button with text "Stop"
      const candidates = [
        '[aria-label="Stop response"]',
        '[aria-label*="Stop"]',
        '[aria-label*="Cancel"]',
        '[data-slot*="cancel"]',
        '[data-slot*="stop"]'
      ]
      for (const sel of candidates) {
        const el = document.querySelector(sel)
        if (el) { el.click(); return { clicked: sel } }
      }
      // Fallback: any button whose textContent includes "Stop"
      for (const b of document.querySelectorAll('button')) {
        if ((b.textContent || '').toLowerCase().includes('stop')) { b.click(); return { clicked: 'btn-text-Stop' } }
      }
      return { clicked: null }
    })()
    `
  )
}

async function main() {
  const tgt = await pickRenderer()
  console.log('target', tgt.url)
  const cdp = await connect(tgt.webSocketDebuggerUrl)
  await cdp.send('Runtime.enable')

  const samples = []
  for (let i = 1; i <= ROUNDS; i++) {
    await focusAndType(cdp, `latency test ${i} ${'x'.repeat(40)}`)
    await new Promise(r => setTimeout(r, 300))
    const result = await submitAndMeasure(cdp, 4000)
    samples.push({ round: i, ...result })
    console.log(
      `r${i}: clear=${(result.composerClearedMs ?? -1).toFixed?.(0) ?? '?'}ms ` +
        `userMsg=${(result.userMessageRenderedMs ?? -1).toFixed?.(0) ?? '?'}ms ` +
        `paint=${(result.userMessagePaintMs ?? -1).toFixed?.(0) ?? '?'}ms ` +
        `reason=${result.reason}`
    )
    // Interrupt the running agent turn so we don't burn tokens on these probes.
    await new Promise(r => setTimeout(r, 200))
    await tryCancel(cdp)
    // Wait long enough for the cancel to settle before next round.
    await new Promise(r => setTimeout(r, 1500))
  }
  writeFileSync('/tmp/hermes-submit-latency.json', JSON.stringify(samples, null, 2))
  console.log('\nwrote /tmp/hermes-submit-latency.json')
  cdp.close()
}

main().catch(e => {
  console.error('fatal:', e.stack ?? e.message)
  process.exit(1)
})