petal/web/src/hooks/useCheckpoint.ts

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

const DEBOUNCE_MS = 4000

// The deterministic rule pack costs nothing to run, so it gets its own short
// debounce: "a apple" should underline while she's still on the sentence, not
// four seconds after she stops typing and then only once the server answers.
const FAST_MS = 250

// A local finding's identity, used to match a provisional card against its
// persisted twin and to remember the ones she's already dealt with. The span
// start is deliberately left out: the same fix at a shifted offset is the same
// fix, and offsets move under her cursor constantly.
export const findingKey = (f: { original: string; replacement: string }) =>
  `${f.original}${f.replacement}`

// provisionalSuggestion dresses a rule-pack finding as a Suggestion so it can
// flow through the highlight/rail machinery unchanged. The `local:` id prefix is
// the only thing that marks it as not-yet-persisted (see resolveServerId); the
// rail renders it exactly like any other card, which is the existing deliberate
// choice for offline findings — see the note on Suggestion.source.
function provisionalSuggestion(docId: string, f: MechanicsFinding): Suggestion {
  return {
    id: `local:${findingKey(f)}`,
    doc_id: docId,
    from_pos: f.from,
    to_pos: f.to,
    original: f.original,
    replacement: f.replacement,
    explanation: f.explanation,
    type: f.type,
    status: 'pending',
    source: 'local',
    created_at: new Date().toISOString(),
  }
}

// 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[]>([])
  // Rule-pack findings rendered before (or instead of) their persisted twins.
  // Non-empty only between the local pass and the mechanics response — or for as
  // long as that response never comes, which is what makes the offline case work.
  const [provisional, setProvisional] = 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)

  // The fast rule-pack timer, the mechanics submit currently in flight (so an
  // accept on a provisional card can wait for its real id), and the findings she
  // has already accepted or dismissed while they were still provisional — the
  // detector has no memory between runs, so without this they'd come straight
  // back. The server keeps the equivalent record for persisted rows.
  const fastRef = useRef<ReturnType<typeof setTimeout>>(undefined)
  const submitRef = useRef<Promise<Suggestion[]> | null>(null)
  const actionedRef = useRef<Set<string>>(new Set())

  // 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]

  // The document text whose rule-pack findings are already persisted server-side,
  // so the checkpoint doesn't re-submit what the fast pass just filed.
  const submittedTextRef = useRef<string | null>(null)

  // runMechanics is the deterministic half of the loop, and the only half that
  // can be instant: detect locally, show the findings *now*, then persist them.
  // The provisional cards are dropped as soon as the server answers — its reply
  // is authoritative, including the suppression of findings she already actioned.
  // If it never answers (offline, server down) they simply stay, which is the
  // whole point: the rule pack needs no network to be right.
  const runMechanics = useCallback(async (text: string) => {
    const id = docIdRef.current
    if (!id) return
    const run = runRef.current
    const findings = mechanicsFindings(text).filter((f) => !actionedRef.current.has(findingKey(f)))
    if (run === runRef.current && id === docIdRef.current) {
      setProvisional(findings.map((f) => provisionalSuggestion(id, f)))
    }
    const submit = api.submitMechanics(id, findings)
    submitRef.current = submit
    try {
      const unified = await submit
      submittedTextRef.current = text
      if (run === runRef.current && id === docIdRef.current) {
        setSuggestions(unified)
        setProvisional([])
      }
    } catch (err) {
      // Best-effort: a mechanics failure leaves the provisional cards standing
      // and must never block the grammar pass.
      console.error('mechanics submit failed', err)
    } finally {
      if (submitRef.current === submit) submitRef.current = null
    }
  }, [])

  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, so the checkpoint's unified response
      // already carries them. Normally the fast pass filed this exact text
      // seconds ago (both timers reset on every edit, so they see the same
      // document) and there's nothing to do; this is the catch-up path for when
      // that submit failed. Only on the initial try — a grammar retry shouldn't
      // re-submit unchanged findings.
      const text = latestTextRef.current
      if (attempt === 0 && text !== null && text !== submittedTextRef.current) {
        await runMechanics(text)
      }
      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)
    }
  }, [runMechanics])

  // 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)
    // The rule pack runs on its own much shorter fuse. Only when the text
    // actually changed: a tone-only schedule has nothing new to detect.
    if (text !== undefined) {
      clearTimeout(fastRef.current)
      fastRef.current = setTimeout(() => void runMechanics(text), FAST_MS)
    }
  }, [runCheck, runMechanics])

  // 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)
    clearTimeout(fastRef.current)
    runRef.current++
    latestTextRef.current = null // unknown until the new doc's first edit
    submittedTextRef.current = null
    submitRef.current = null
    actionedRef.current.clear()
    setProvisional([])
    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)
      clearTimeout(fastRef.current)
    },
    [],
  )

  // Drop one suggestion locally (after accept/dismiss) without a refetch. A
  // provisional card is also remembered as actioned, so the next local pass —
  // which reads only the text, and may still see the same wording if she
  // dismissed rather than accepted — doesn't hand it straight back.
  const removeSuggestion = useCallback((id: string) => {
    setSuggestions((prev) => prev.filter((s) => s.id !== id))
    setProvisional((prev) => {
      for (const s of prev) if (s.id === id) actionedRef.current.add(findingKey(s))
      return prev.filter((s) => s.id !== id)
    })
  }, [])

  // resolveServerId maps a card to the row the API can act on. Persisted cards
  // are themselves; a provisional one waits for its submit to land and then
  // finds its twin, so accepting inside that window still records the keep (and
  // plants its word in the garden) instead of quietly dropping it. Null means
  // there is no row — offline, or the server suppressed the finding — in which
  // case the caller has already applied the edit and there's nothing to file.
  const resolveServerId = useCallback(async (s: Suggestion): Promise<string | null> => {
    if (!s.id.startsWith('local:')) return s.id
    const pending = submitRef.current
    if (!pending) return null
    try {
      const unified = await pending
      return unified.find((x) => findingKey(x) === findingKey(s))?.id ?? null
    } catch {
      return null
    }
  }, [])

  // What the editor sees: the server's set, plus any rule-pack finding not yet
  // represented in it. Matching on wording rather than position keeps a card
  // from flickering into a duplicate while she types around it.
  const merged = useMemo(() => {
    if (provisional.length === 0) return suggestions
    const known = new Set(suggestions.map(findingKey))
    return [...suggestions, ...provisional.filter((p) => !known.has(findingKey(p)))]
  }, [suggestions, provisional])

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