ultisuite-client/lib/drive/extensions/docs-page-flow-decoration.ts

import { Extension } from "@tiptap/core"
import type { Editor } from "@tiptap/react"
import { Plugin, PluginKey } from "@tiptap/pm/state"
import { Decoration, DecorationSet, type EditorView } from "@tiptap/pm/view"

export const PAGE_FLOW_PLUGIN_KEY = new PluginKey<DecorationSet>("docsPageFlowDecoration")

function decorationSetsEqual(
  current: DecorationSet | undefined,
  next: DecorationSet
): boolean {
  if (!current || current === DecorationSet.empty) {
    return next === DecorationSet.empty || next.find(0, Number.MAX_SAFE_INTEGER).length === 0
  }
  const a = current.find(0, Number.MAX_SAFE_INTEGER)
  const b = next.find(0, Number.MAX_SAFE_INTEGER)
  if (a.length !== b.length) return false
  for (let i = 0; i < a.length; i++) {
    if (a[i].from !== b[i].from) return false
    const aPush = (a[i].spec as { pushPx?: number }).pushPx
    const bPush = (b[i].spec as { pushPx?: number }).pushPx
    if (aPush !== bPush) return false
  }
  return true
}

export function readPageFlowMetrics(prose: HTMLElement): {
  bodyAreaH: number
  interPageSpacer: number
} | null {
  const surface = prose.closest(".ultidrive-docs-editor-surface")
  if (!surface) return null
  const styles = getComputedStyle(surface)
  const bodyAreaH = parseFloat(styles.getPropertyValue("--docs-body-area-h"))
  const interPageSpacer = parseFloat(styles.getPropertyValue("--docs-inter-page-spacer"))
  if (!Number.isFinite(bodyAreaH) || bodyAreaH <= 0) return null
  if (!Number.isFinite(interPageSpacer) || interPageSpacer <= 0) return null
  return { bodyAreaH, interPageSpacer }
}

function blockDomAtOffset(view: EditorView, offset: number): HTMLElement | null {
  const dom = view.nodeDOM(offset)
  if (dom instanceof HTMLElement) {
    if (dom.parentElement === view.dom) return dom
    let el: HTMLElement | null = dom
    while (el && el.parentElement !== view.dom) el = el.parentElement
    return el
  }
  return null
}

/** Block height in prose px (offsetHeight + bottom margin; immune to canvas scale). */
function measureBlockFlowHeight(dom: HTMLElement): number {
  const style = getComputedStyle(dom)
  const marginBottom = parseFloat(style.marginBottom) || 0
  return dom.offsetHeight + marginBottom
}

function createSpacerElement(height: number): HTMLElement {
  const el = document.createElement("div")
  el.className = "docs-page-flow-spacer"
  el.style.cssText = `display:block;width:100%;height:${height}px;margin:0;padding:0;border:0;flex-shrink:0`
  el.setAttribute("aria-hidden", "true")
  el.contentEditable = "false"
  return el
}

/** Simulate vertical flow without reading pushed DOM positions (avoids clearing decorations). */
export function computePageFlowPushes(
  blocks: Array<{ height: number }>,
  bodyAreaH: number,
  interPageSpacer: number
): Array<{ blockIndex: number; pushPx: number; breakY: number }> {
  const pageStep = bodyAreaH + interPageSpacer
  const pushes: Array<{ blockIndex: number; pushPx: number; breakY: number }> = []
  let simulatedY = 0
  let nextBreakY = bodyAreaH

  blocks.forEach(({ height: blockH }, blockIndex) => {
    if (blockH <= 0) return

    while (simulatedY + blockH > nextBreakY) {
      if (simulatedY < nextBreakY) {
        const pushPx = nextBreakY - simulatedY + interPageSpacer
        pushes.push({ blockIndex, pushPx, breakY: nextBreakY })
        simulatedY = nextBreakY + interPageSpacer
        nextBreakY += pageStep
        break
      }
      nextBreakY += pageStep
    }

    simulatedY += blockH
  })

  return pushes
}

export function computeSimulatedLayoutHeight(
  blocks: Array<{ height: number }>,
  bodyAreaH: number,
  interPageSpacer: number
): number {
  const pageStep = bodyAreaH + interPageSpacer
  let simulatedY = 0
  let nextBreakY = bodyAreaH

  for (const { height: blockH } of blocks) {
    if (blockH <= 0) continue
    while (simulatedY + blockH > nextBreakY) {
      if (simulatedY < nextBreakY) {
        simulatedY = nextBreakY + interPageSpacer
        nextBreakY += pageStep
        break
      }
      nextBreakY += pageStep
    }
    simulatedY += blockH
  }

  return simulatedY
}

function collectFlowBlocks(view: EditorView): Array<{ height: number; offset: number }> {
  const blocks: Array<{ height: number; offset: number }> = []
  view.state.doc.forEach((node, offset) => {
    if (!node.isBlock || node.type.name === "doc") return
    const dom = blockDomAtOffset(view, offset)
    if (!dom) return
    const blockH = measureBlockFlowHeight(dom)
    if (blockH <= 0) return
    blocks.push({ height: blockH, offset })
  })
  return blocks
}

/** Build page-flow spacer widgets for blocks that cross a page body boundary. */
export function buildPageFlowDecorations(view: EditorView): DecorationSet {
  const metrics = readPageFlowMetrics(view.dom)
  if (!metrics) return DecorationSet.empty

  const { bodyAreaH, interPageSpacer } = metrics
  const blocks = collectFlowBlocks(view)

  const pushes = computePageFlowPushes(
    blocks.map((b) => ({ height: b.height })),
    bodyAreaH,
    interPageSpacer
  )

  const decorations = pushes.map(({ blockIndex, pushPx, breakY }) => {
    const block = blocks[blockIndex]
    return Decoration.widget(
      block.offset,
      () => createSpacerElement(pushPx),
      { side: -1, key: `page-flow-spacer-${block.offset}-${breakY}`, pushPx }
    )
  })

  return DecorationSet.create(view.state.doc, decorations)
}

/** Apply page-flow layout once. Returns true if decorations changed. */
export function applyPageFlowLayout(editor: Editor): boolean {
  if (editor.isDestroyed || !editor.isInitialized) return false
  const view = editor.view
  const next = buildPageFlowDecorations(view)
  const current = PAGE_FLOW_PLUGIN_KEY.getState(view.state)
  if (decorationSetsEqual(current, next)) return false
  const tr = view.state.tr.setMeta(PAGE_FLOW_PLUGIN_KEY, next)
  tr.setMeta("addToHistory", false)
  view.dispatch(tr)
  return true
}

/** Pure helper for tests: return count of inter-page pushes. */
export function countPageFlowSpacers(
  blocks: Array<{ height: number }>,
  bodyAreaH: number,
  interPageSpacer: number
): number {
  return computePageFlowPushes(blocks, bodyAreaH, interPageSpacer).length
}

/** State-only plugin: layout is driven from docs-page-view (no view plugin / observers). */
export const DocsPageFlowDecoration = Extension.create({
  name: "docsPageFlowDecoration",

  addProseMirrorPlugins() {
    return [
      new Plugin({
        key: PAGE_FLOW_PLUGIN_KEY,
        state: {
          init: () => DecorationSet.empty,
          apply(tr, set) {
            const next = tr.getMeta(PAGE_FLOW_PLUGIN_KEY)
            if (next instanceof DecorationSet) return next
            return set.map(tr.mapping, tr.doc)
          },
        },
        props: {
          decorations(state) {
            return PAGE_FLOW_PLUGIN_KEY.getState(state) ?? DecorationSet.empty
          },
        },
      }),
    ]
  },
})