petal/web/src/hooks/useCheckpoint.ts

import { useCallback, useEffect, useRef, useState } from 'react'
import { api, type Suggestion } from '../api/client'
import { mechanicsFindings } from '../components/Companion/prose'

const DEBOUNCE_MS = 4000

// useCheckpoint manages the grammar-checkpoint lifecycle for one document:
// it loads any existing pending suggestions when the doc opens, then fires a
// fresh check 4s after the user stops typing. `checking` drives the breathing
// dot in the StatusBar. The server rate-limits per document, so a check that
// fires too soon simply returns the current set unchanged.
export function useCheckpoint(docId: string | null) {
  const [suggestions, setSuggestions] = useState<Suggestion[]>([])
  const [checking, setChecking] = useState(false)
  // True while a whole-document voice pass is in flight (explicit user action).
  const [voicing, setVoicing] = useState(false)
  // True while a whole-document collocation pass is in flight (explicit action).
  const [collocating, setCollocating] = useState(false)
  // True when the last LLM pass couldn't reach the model (server 502 or network
  // error). Drives a gentle, reassuring "helper is resting" note — the writing
  // itself still saves fine, so this is awareness, not an error. Cleared on the
  // next success or doc switch.
  const [llmDown, setLlmDown] = useState(false)

  const debounceRef = useRef<ReturnType<typeof setTimeout>>(undefined)
  // Pending auto-retry timer for a failed checkpoint (see runCheck).
  const retryRef = useRef<ReturnType<typeof setTimeout>>(undefined)
  const docIdRef = useRef(docId)
  docIdRef.current = docId

  // Latest plaintext, kept fresh by schedule(), so the deterministic mechanics
  // pass (detected client-side, see prose.ts) reads the current document when the
  // debounce fires — no need to re-thread text through every call. null until the
  // first edit of the current doc, so a tone-only check before any edit doesn't
  // submit an empty batch and wipe the doc's persisted mechanics rows.
  const latestTextRef = useRef<string | null>(null)

  // Token to discard responses from a doc we've since navigated away from.
  const runRef = useRef(0)

  // Backoff schedule for re-running a checkpoint that failed. A paste fires
  // exactly ONE checkpoint, and nothing re-fires until the next keystroke — so a
  // single transient failure (a momentary stall on the shared inference box)
  // would otherwise strand the user on "resting" indefinitely after a paste.
  // These delays clear the server's 30s per-doc floor by the later attempts, and
  // a failed pass now releases its slot server-side so a retry can truly re-run.
  const RETRY_DELAYS_MS = [3000, 12000, 35000]

  const runCheck = useCallback(async (attempt = 0) => {
    const id = docIdRef.current
    if (!id) return
    clearTimeout(retryRef.current)
    const run = ++runRef.current
    setChecking(true)
    let retrying = false
    try {
      // Deterministic mechanics first: detect client-side and persist as the
      // 'mechanics' family before the grammar pass, so the checkpoint's unified
      // response already carries them. Best-effort and only on the initial try —
      // a grammar retry shouldn't re-submit unchanged findings. A mechanics
      // failure must not block the grammar pass.
      if (attempt === 0 && latestTextRef.current !== null) {
        try {
          // Render the mechanics fixes immediately — they're instant (no LLM), so
          // the unified set they return shouldn't wait on the slow grammar pass.
          const withMech = await api.submitMechanics(id, mechanicsFindings(latestTextRef.current))
          if (run === runRef.current && id === docIdRef.current) setSuggestions(withMech)
        } catch (err) {
          console.error('mechanics submit failed', err)
        }
      }
      const fresh = await api.checkDoc(id)
      if (run === runRef.current && id === docIdRef.current) {
        setSuggestions(fresh)
        setLlmDown(false)
      }
    } catch (err) {
      console.error('checkpoint failed', err)
      if (run !== runRef.current) return
      if (attempt < RETRY_DELAYS_MS.length) {
        // Keep trying quietly — don't flag "resting" until retries are exhausted.
        retrying = true
        retryRef.current = setTimeout(() => void runCheck(attempt + 1), RETRY_DELAYS_MS[attempt])
      } else {
        setLlmDown(true)
      }
    } finally {
      // Stay in the "checking" state while a retry is queued so the breathing dot
      // keeps reassuring rather than flickering off between attempts.
      if (run === runRef.current && !retrying) setChecking(false)
    }
  }, [])

  // runExplicitPass drives a whole-document, explicit-action pass (voice,
  // collocation): it returns the unified pending set, so the other families'
  // highlights survive, and shares the run token so navigating away discards a
  // late response. It supersedes any queued grammar retry AND clears the
  // checkpoint's "checking" state, so the breathing rose dot can't linger on
  // after the retry timer it would have cleared is cancelled here.
  const runExplicitPass = useCallback(
    async (call: (id: string) => Promise<Suggestion[]>, setBusy: (b: boolean) => void, label: string) => {
      const id = docIdRef.current
      if (!id) return
      // An explicit action fully supersedes a queued auto-check and grammar
      // retry, and takes over the indicator — clear the stranded "checking" dot.
      clearTimeout(debounceRef.current)
      clearTimeout(retryRef.current)
      setChecking(false)
      const run = ++runRef.current
      setBusy(true)
      try {
        const full = await call(id)
        if (run === runRef.current && id === docIdRef.current) {
          setSuggestions(full)
          setLlmDown(false)
        }
      } catch (err) {
        console.error(`${label} failed`, err)
        if (run === runRef.current) setLlmDown(true)
      } finally {
        // setBusy is THIS invocation's own flag (a newer run sets its own), so
        // clear it unconditionally — otherwise an overlapping pass that bumped
        // the run token would strand this spinner on "Reading…" forever.
        setBusy(false)
      }
    },
    [],
  )

  // Run the voice-consistency pass now (explicit "Check my voice" action).
  const runVoice = useCallback(
    () => runExplicitPass(api.voiceDoc, setVoicing, 'voice pass'),
    [runExplicitPass],
  )

  // Run the collocation coach now (explicit "Make it sound natural" action).
  const runCollocation = useCallback(
    () => runExplicitPass(api.collocationDoc, setCollocating, 'collocation pass'),
    [runExplicitPass],
  )

  // Call on every edit; schedules a check 4s after typing settles. Pass the
  // current plaintext so the mechanics pass sees the latest document; omit it
  // (e.g. a tone-only change) to reuse the last text.
  const schedule = useCallback((text?: string) => {
    if (text !== undefined) latestTextRef.current = text
    clearTimeout(debounceRef.current)
    clearTimeout(retryRef.current) // a fresh edit supersedes any queued retry
    debounceRef.current = setTimeout(() => void runCheck(), DEBOUNCE_MS)
  }, [runCheck])

  // Load existing pending suggestions whenever the document changes, and cancel
  // any in-flight debounce from the previous doc.
  useEffect(() => {
    clearTimeout(debounceRef.current)
    clearTimeout(retryRef.current)
    runRef.current++
    latestTextRef.current = null // unknown until the new doc's first edit
    setSuggestions([])
    setChecking(false)
    setVoicing(false)
    setCollocating(false)
    setLlmDown(false)
    if (!docId) return
    let cancelled = false
    void (async () => {
      try {
        const existing = await api.listSuggestions(docId)
        if (!cancelled && docIdRef.current === docId) setSuggestions(existing)
      } catch (err) {
        console.error('failed to load suggestions', err)
      }
    })()
    return () => {
      cancelled = true
    }
  }, [docId])

  useEffect(
    () => () => {
      clearTimeout(debounceRef.current)
      clearTimeout(retryRef.current)
    },
    [],
  )

  // Drop one suggestion locally (after accept/dismiss) without a refetch.
  const removeSuggestion = useCallback((id: string) => {
    setSuggestions((prev) => prev.filter((s) => s.id !== id))
  }, [])

  return { suggestions, checking, voicing, collocating, llmDown, schedule, runVoice, runCollocation, removeSuggestion }
}