write/frontend/src/components/Editor.jsx

import { forwardRef, useEffect, useRef, useImperativeHandle, useState, useCallback } from 'react'
import { createPortal } from 'react-dom'
import { useEditor, EditorContent, ReactNodeViewRenderer } from '@tiptap/react'
import { Mark, mergeAttributes, getMarkRange } from '@tiptap/core'
import StarterKit from '@tiptap/starter-kit'
import Image from '@tiptap/extension-image'
import Placeholder from '@tiptap/extension-placeholder'
import Underline from '@tiptap/extension-underline'
import CharacterCount from '@tiptap/extension-character-count'
import TextAlign from '@tiptap/extension-text-align'
import Highlight from '@tiptap/extension-highlight'
import Toolbar from './Toolbar'
import ImageView from './ImageView'
import { api } from '../lib/api'

// ── Custom image extension ─────────────────────────────────────────────────
const CustomImage = Image.extend({
  addAttributes() {
    return {
      ...this.parent?.(),
      width: { default: '100%' },
      align: { default: 'center' },
    }
  },
  addNodeView() {
    return ReactNodeViewRenderer(ImageView)
  },
}).configure({ allowBase64: false, inline: false })

// ── ReadingMark — transient mark that follows the TTS playhead ────────────
const ReadingMark = Mark.create({
  name: 'readingWord',
  inclusive: false,
  parseHTML() { return [] },   // never restore from HTML — transient UI only
  renderHTML({ HTMLAttributes }) {
    return ['span', mergeAttributes({ class: 'reading-word' }, HTMLAttributes), 0]
  },
})

// ── LintMark — transient mark for spelling / grammar underlines ────────────
const LintMark = Mark.create({
  name: 'lintError',
  inclusive: false,   // don't extend when typing at a boundary

  addAttributes() {
    return {
      kind:    { default: 'grammar' },   // 'spelling' | 'grammar'
      message: { default: '' },
      replacements: {
        default: [],
        parseHTML:  el  => JSON.parse(el.getAttribute('data-r') || '[]'),
        renderHTML: att => ({ 'data-r': JSON.stringify(att.replacements ?? []) }),
      },
    }
  },

  parseHTML() { return [] },  // never parse from saved HTML — marks are transient

  renderHTML({ HTMLAttributes }) {
    return ['span', mergeAttributes(
      { 'data-lint': '', class: `lint-mark lint-mark--${HTMLAttributes.kind}` },
      HTMLAttributes,
    ), 0]
  },
})

// ── Lint helpers ───────────────────────────────────────────────────────────

// Build plain text + ProseMirror position map from the doc
function buildTextMap(doc) {
  const text = [], posMap = []
  doc.descendants((node, pos) => {
    if (node.isText) {
      for (let i = 0; i < node.text.length; i++) {
        text.push(node.text[i])
        posMap.push(pos + i)
      }
      return false
    }
    // Paragraph / heading separators → '\n' (no PM position — null in map)
    if (node.isBlock && text.length > 0 && text[text.length - 1] !== '\n') {
      text.push('\n')
      posMap.push(null)
    }
  })
  return { text: text.join(''), posMap }
}

// Remove all lint marks in one transaction
function clearLintMarks(editor) {
  const { tr, doc, schema } = editor.state
  const mt = schema.marks.lintError
  if (!mt) return
  doc.descendants((node, pos) => {
    if (!node.isText) return
    node.marks.filter(m => m.type === mt)
      .forEach(() => tr.removeMark(pos, pos + node.nodeSize, mt))
  })
  editor.view.dispatch(tr)
}

// Strip transient UI marks (lint, reading playhead) from JSON before saving
function stripLintMarks(node) {
  if (!node) return node
  const n = { ...node }
  if (n.marks) {
    n.marks = n.marks.filter(m => m.type !== 'lintError' && m.type !== 'readingWord')
    if (!n.marks.length) delete n.marks
  }
  if (n.content) n.content = n.content.map(stripLintMarks)
  return n
}

// ── Editor component ───────────────────────────────────────────────────────
const Editor = forwardRef(function Editor(
  { content, onChange, onImageUpload, fontSize, onFontSizeChange }, ref,
) {
  const synced = useRef(false)

  // ── Lint state ───────────────────────────────────────────────────────────
  const [lintStatus, setLintStatus]     = useState('idle')   // idle | checking | done | stale
  const [lintCount,  setLintCount]      = useState(0)
  const [lintPopover, setLintPopover]   = useState(null)     // { top, left, message, replacements, from, to }
  const lintStatusRef  = useRef('idle')
  const lintDebounce   = useRef(null)
  const runLintRef     = useRef(null)   // always holds latest runLint, safe to call from onUpdate closure
  const applyingLints  = useRef(false)  // true while dispatching mark transactions — suppresses onUpdate

  const editor = useEditor({
    extensions: [
      StarterKit,
      Underline,
      CustomImage,
      LintMark,
      ReadingMark,
      Placeholder.configure({ placeholder: 'Begin your story here…' }),
      CharacterCount,
      TextAlign.configure({ types: ['heading', 'paragraph'] }),
      Highlight.configure({ multicolor: true }),
    ],
    content: '',
    onUpdate({ editor }) {
      // Skip when the update was caused by us applying / clearing lint marks.
      // Without this guard, dispatching mark transactions fires onUpdate which
      // reschedules the lint check, which applies marks, which fires onUpdate… ∞
      if (applyingLints.current) return
      onChange(stripLintMarks(editor.getJSON()))
      // Debounced auto-check: fire 1.5 s after the last keystroke
      clearTimeout(lintDebounce.current)
      lintDebounce.current = setTimeout(() => runLintRef.current?.(), 3000)
    },
  })

  useImperativeHandle(ref, () => ({
    insertPrompt: (text) => {
      if (!editor) return
      editor.chain().focus().insertContent({
        type: 'blockquote',
        content: [{ type: 'paragraph', content: [{ type: 'text', text }] }],
      }).run()
    },
  }))

  useEffect(() => {
    if (editor && !synced.current) {
      const hasContent = content && Object.keys(content).length > 0
      if (hasContent) editor.commands.setContent(content, false)
      synced.current = true
    }
  }, [editor, content])

  // ── Lint check ───────────────────────────────────────────────────────────
  const runLint = useCallback(async () => {
    if (!editor || lintStatusRef.current === 'checking') return
    lintStatusRef.current = 'checking'
    setLintStatus('checking')
    // Don't close an open popover here — it gets wiped on every auto-check
    // and the user can't interact with it. It closes naturally on click-outside.

    try {
      const { text, posMap } = buildTextMap(editor.state.doc)
      if (!text.trim()) {
        applyingLints.current = true
        clearLintMarks(editor)
        applyingLints.current = false
        setLintCount(0)
        setLintStatus('done')
        lintStatusRef.current = 'done'
        return
      }

      const data = await api.lintCheck(text)

      // Clear old marks then apply new ones — both wrapped in the flag so
      // the resulting onUpdate calls are silently ignored.
      applyingLints.current = true
      clearLintMarks(editor)
      const { tr, schema } = editor.state
      const lintType = schema.marks.lintError

      for (const match of data.matches) {
        let pmFrom = null, pmTo = null
        for (let i = match.offset; i < match.offset + match.length && i < posMap.length; i++) {
          if (posMap[i] !== null) {
            if (pmFrom === null) pmFrom = posMap[i]
            pmTo = posMap[i] + 1
          }
        }
        if (pmFrom === null) continue
        tr.addMark(pmFrom, pmTo, lintType.create({
          kind:         match.kind,
          message:      match.message,
          replacements: match.replacements,
        }))
      }
      editor.view.dispatch(tr)
      applyingLints.current = false

      setLintCount(data.matches.length)
      setLintStatus('done')
      lintStatusRef.current = 'done'
    } catch (err) {
      applyingLints.current = false
      console.error('[lint] Check failed:', err)
      setLintStatus('idle')
      lintStatusRef.current = 'idle'
    }
  }, [editor])

  // Keep ref in sync so the stale onUpdate closure always calls the latest runLint
  useEffect(() => { runLintRef.current = runLint }, [runLint])
  // Clean up the debounce timer on unmount
  useEffect(() => () => clearTimeout(lintDebounce.current), [])

  // ── Reading-word mark (TTS playhead highlight) ────────────────────────────
  // Called by Toolbar on each word-boundary event.  Both operations are wrapped
  // in applyingLints so the resulting onUpdate dispatch is silently ignored —
  // the same guard that protects lint mark dispatches also covers these.
  const applyReadingMark = useCallback((from, to) => {
    if (!editor) return
    const mt = editor.state.schema.marks.readingWord
    if (!mt) return
    const { tr } = editor.state
    // Clear any existing reading mark in one shot then set the new one
    editor.state.doc.descendants((node, pos) => {
      if (!node.isText) return
      node.marks.filter(m => m.type === mt)
        .forEach(() => tr.removeMark(pos, pos + node.nodeSize, mt))
    })
    tr.addMark(from, to, mt.create())
    applyingLints.current = true
    editor.view.dispatch(tr)
    applyingLints.current = false
  }, [editor])

  const clearReadingMark = useCallback(() => {
    if (!editor) return
    const mt = editor.state.schema.marks.readingWord
    if (!mt) return
    const { tr } = editor.state
    editor.state.doc.descendants((node, pos) => {
      if (!node.isText) return
      node.marks.filter(m => m.type === mt)
        .forEach(() => tr.removeMark(pos, pos + node.nodeSize, mt))
    })
    applyingLints.current = true
    editor.view.dispatch(tr)
    applyingLints.current = false
  }, [editor])

  // ── Popover on clicking a lint mark ──────────────────────────────────────
  function handleEditorClick(e) {
    if (!editor) return
    const target = e.target.closest('[data-lint]')
    if (!target) { setLintPopover(null); return }

    const coords = editor.view.posAtCoords({ left: e.clientX, top: e.clientY })
    if (!coords) return

    const $pos    = editor.state.doc.resolve(coords.pos)
    const lintMT  = editor.state.schema.marks.lintError
    const mark    = $pos.marks().find(m => m.type === lintMT)
    if (!mark) return

    const range = getMarkRange($pos, lintMT)
    if (!range) return

    const rect = target.getBoundingClientRect()
    setLintPopover({
      // Fixed position so we don't have to fight scroll offsets
      top:  rect.bottom + 8,
      left: Math.max(8, Math.min(rect.left, window.innerWidth - 270)),
      message:      mark.attrs.message,
      replacements: mark.attrs.replacements || [],
      from:         range.from,
      to:           range.to,
    })
  }

  function applyReplacement(replacement) {
    if (!lintPopover || !editor) return
    const { from, to } = lintPopover
    editor.chain().setTextSelection({ from, to }).insertContent(replacement).focus().run()
    setLintPopover(null)
    // Re-run lint after a short pause so the new text gets checked
    setTimeout(runLint, 600)
  }

  function dismissLintMark() {
    if (!lintPopover || !editor) return
    const { from, to } = lintPopover
    const { tr, schema } = editor.state
    tr.removeMark(from, to, schema.marks.lintError)
    editor.view.dispatch(tr)
    setLintPopover(null)
    setLintCount(c => Math.max(0, c - 1))
  }

  const wordCount = editor?.storage.characterCount.words() ?? 0

  return (
    <div className="editor-outer" style={{ '--editor-font-size': (fontSize || 17) + 'px' }}>
      <Toolbar
        editor={editor}
        onImageUpload={onImageUpload}
        fontSize={fontSize}
        onFontSizeChange={onFontSizeChange}
        onLint={runLint}
        lintStatus={lintStatus}
        lintCount={lintCount}
        applyReadingMark={applyReadingMark}
        clearReadingMark={clearReadingMark}
      />
      <div className="editor-wrap" onClick={handleEditorClick}>
        <EditorContent editor={editor} className="editor-body" spellCheck={false} />
        <div className="word-count">{wordCount.toLocaleString()} {wordCount === 1 ? 'word' : 'words'}</div>
      </div>

      {lintPopover && createPortal(
        <div
          className="lint-popover"
          style={{ top: lintPopover.top, left: lintPopover.left }}
          // Prevent mousedown from stealing focus away from the editor —
          // without this the editor blurs and the popover may vanish before click fires
          onMouseDown={e => e.preventDefault()}
          onClick={e => e.stopPropagation()}
        >
          <p className="lint-popover-msg">{lintPopover.message}</p>
          {lintPopover.replacements.length > 0 && (
            <div className="lint-suggestions">
              {lintPopover.replacements.map((r, i) => (
                <button key={i} className="lint-suggestion" onClick={() => applyReplacement(r)}>{r}</button>
              ))}
            </div>
          )}
          <div className="lint-popover-actions">
            <button className="btn btn-ghost lint-ignore" onClick={dismissLintMark}>Ignore</button>
            <button className="btn btn-ghost lint-close"  onClick={() => setLintPopover(null)}>✕</button>
          </div>
        </div>,
        document.body,
      )}
    </div>
  )
})

export default Editor