studio/src/hooks/useStream.js

import useAppStore from '../store/appStore'
import { getTree, readFile } from '../api/projects'

// Sentinel kept exported for the v1 retry indicator (used by historical messages).
const RETRYING_MARKER = '__TESSERACT_RETRYING__'

export function useStream() {
  async function sendMessage(text, { pinnedFiles = [] } = {}) {
    const state = useAppStore.getState()
    const { activeSessionId, project } = state
    if (!activeSessionId || !text.trim()) return

    state.appendMessage({ role: 'user', content: text })
    state.appendMessage({ role: 'assistant', content: '', toolCalls: [] })
    state.setIsStreaming(true)
    // Per-turn metrics (tokens/s, elapsed, last-turn cost) reset to clean
    // values; sessionTokens is the cumulative number and is NOT reset here.
    state.resetInferenceStats()
    state.setInferenceStats({ visible: true })

    const startTime = performance.now()

    try {
      await _stream({
        sessionId:    activeSessionId,
        message:      text,
        activeFile:   project.activeFile,
        pinnedFiles,
        startTime,
      })

      // Refresh the project tree + the editor cache for any files the agent
      // touched.  Cheap: a single tree call + selective re-reads for tracked
      // paths that already existed (so Monaco's in-memory copy stays consistent).
      if (project.id) {
        await _resyncProject(project.id)
      }
      // no-op
    } catch (err) {
      if (err?.name !== 'AbortError') {
        console.error('Stream error:', err)
        useAppStore.getState().setLastAssistantError(
          'Something went wrong while talking to the model. Please try again.',
        )
      }
    } finally {
      useAppStore.getState().setIsStreaming(false)
      useAppStore.getState().setAbortController(null)
      useAppStore.getState().clearPacing()
    }
  }

  return { sendMessage }
}

// ── Project re-sync after a turn ────────────────────────────────────────────
// The agent may have written, deleted, or created files.  Refetch the tree;
// for files already in the cache, re-read them so Monaco shows the new content
// when the user opens that tab.
async function _resyncProject(projectId) {
  const store = useAppStore.getState()
  const beforeFiles = store.project.files

  try {
    const tree = await getTree(projectId)
    const treePaths = new Set(tree.filter((e) => !e.is_dir).map((e) => e.path))

    // Drop cached files that no longer exist on disk.
    const nextFiles = {}
    for (const [p, content] of Object.entries(beforeFiles)) {
      if (treePaths.has(p)) nextFiles[p] = content
    }

    // Re-read files we already had open OR the active file — they're the ones
    // the user might be looking at right now.  Lazily re-read others on click.
    const toReread = new Set(Object.keys(nextFiles))
    if (store.project.activeFile) toReread.add(store.project.activeFile)

    const fetched = await Promise.all(
      Array.from(toReread).map((p) =>
        readFile(projectId, p).then((c) => [p, c]).catch(() => null),
      ),
    )
    for (const entry of fetched) {
      if (entry) nextFiles[entry[0]] = entry[1]
    }

    store.setProject({ tree, files: nextFiles })
  } catch (e) {
    console.warn('project resync failed', e)
  }
}

// ── Single SSE stream pass ──────────────────────────────────────────────────
async function _stream({ sessionId, message, activeFile, pinnedFiles, startTime }) {
  const controller = new AbortController()
  useAppStore.getState().setAbortController(controller)

  let firstTokenTime = 0
  let tokenCount = 0
  let lastStatsUpdate = performance.now()
  let serverError = null
  let turnUsage = null         // backend's authoritative usage if it arrives
  let liveSessionAdded = 0     // how much we've optimistically added to sessionTokens
                                // so we can reconcile against the real number at end

  const res = await fetch('/api/chat', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({
      session_id:   sessionId,
      message,
      active_file:  activeFile,
      pinned_files: pinnedFiles?.length ? pinnedFiles : undefined,
    }),
    signal: controller.signal,
  })

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

  while (!done) {
    let value, streamDone
    try {
      ;({ value, done: streamDone } = await reader.read())
    } catch (err) {
      if (err?.name === 'AbortError') return { aborted: true, error: null }
      throw err
    }
    if (streamDone) break

    buffer += decoder.decode(value, { stream: true })

    // SSE messages are separated by a blank line.  Process complete messages
    // out of the buffer; the trailing partial chunk waits for the next read.
    let nl
    while ((nl = buffer.indexOf('\n\n')) !== -1) {
      const block = buffer.slice(0, nl)
      buffer = buffer.slice(nl + 2)

      for (const line of block.split('\n')) {
        if (!line.startsWith('data: ')) continue
        const payload = line.slice(6).trim()
        if (payload === '[DONE]') { done = true; break }

        let event
        try { event = JSON.parse(payload) } catch { continue }

        if (event.error) {
          serverError = event.error
          useAppStore.getState().setLastAssistantError(event.error)
          continue
        }

        if (event.tool) {
          useAppStore.getState().clearPacing()
          useAppStore.getState().upsertLastAssistantTool(event.tool)
          continue
        }

        if (event.usage) {
          // Cumulative usage for the whole turn (prompt + completion across
          // every Groq round).  Source of truth — replaces the local count.
          turnUsage = event.usage
          continue
        }

        if (event.session_renamed) {
          const { id, name } = event.session_renamed
          const s = useAppStore.getState()
          s.setSessions(
            s.sessions.map((x) => (x.id === id ? { ...x, name } : x))
          )
          continue
        }

        if (event.pacing) {
          // Backend is sleeping for the TPM window to roll — surface as a
          // pacing pill, never as an error.
          const { wait_seconds, reason, snapshot } = event.pacing
          useAppStore.getState().setPacing({
            waitSeconds: wait_seconds,
            reason,
            since:       performance.now(),
          })
          if (snapshot) useAppStore.getState().setRateLimit(snapshot)
          continue
        }

        if (event.token) {
          if (firstTokenTime === 0) {
            firstTokenTime = performance.now()
            useAppStore.getState().setInferenceStats({
              firstTokenMs: Math.round(firstTokenTime - startTime),
            })
          }
          useAppStore.getState().clearPacing()
          useAppStore.getState().updateLastAssistantMessage(event.token)
          tokenCount++

          // Optimistic live increment so the session-tok counter ticks up as
          // the user watches text stream in.  Reconciled at end of turn
          // against the backend's authoritative usage number.
          useAppStore.getState().addSessionTokens(1)
          liveSessionAdded += 1

          const now = performance.now()
          if (now - lastStatsUpdate > 150) {
            const generationMs = now - firstTokenTime
            const tps = generationMs > 0 ? tokenCount / (generationMs / 1000) : 0
            useAppStore.getState().setInferenceStats({
              tokens:       tokenCount,
              tokensPerSec: Math.round(tps * 10) / 10,
              elapsedMs:    Math.round(now - startTime),
            })
            lastStatsUpdate = now
          }
        }
      }
    }
  }

  const endTime = performance.now()
  const elapsedMs = endTime - startTime
  const generationMs = firstTokenTime > 0 ? endTime - firstTokenTime : elapsedMs
  const tps = generationMs > 0 ? tokenCount / (generationMs / 1000) : 0

  // Prefer the real Groq usage when it arrived; fall back to the local count
  // only if the turn errored before usage was emitted.
  const realTotal  = turnUsage?.total_tokens ?? 0
  const realPrompt = turnUsage?.prompt_tokens ?? 0
  const realOutput = turnUsage?.completion_tokens ?? 0

  useAppStore.getState().setInferenceStats({
    tokens:       realOutput || tokenCount,           // output tokens drive tok/s
    tokensPerSec: Math.round(tps * 10) / 10,
    elapsedMs:    Math.round(elapsedMs),
    promptTokens: realPrompt,
    totalTokens:  realTotal,
  })

  // Reconcile: we've optimistically added `liveSessionAdded` to sessionTokens
  // during streaming.  Replace that with the real Groq-billed total (prompt +
  // completion, includes everything the API saw).  If usage never arrived,
  // keep the live count — it's at least non-zero.
  if (realTotal > 0) {
    const delta = realTotal - liveSessionAdded
    if (delta !== 0) useAppStore.getState().addSessionTokens(delta)
  }

  return { aborted: false, error: serverError }
}

export { RETRYING_MARKER }