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

import { Extension } from "@tiptap/core"
import type { Editor } from "@tiptap/react"
import type { Node as PMNode } from "@tiptap/pm/model"
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")

type FlowUnit = { height: number; offset: number }

const LIST_TYPES = new Set(["bulletList", "orderedList", "taskList"])

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
}

function nestedDomAt(view: EditorView, offset: number, selector: string): HTMLElement | null {
  const dom = view.nodeDOM(offset)
  if (!(dom instanceof HTMLElement)) return null
  if (dom.matches(selector)) return dom
  return dom.closest(selector)
}

/** 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 measureVisualLineHeights(dom: HTMLElement): Array<{ top: number; height: number }> {
  const blockTop = dom.getBoundingClientRect().top
  const range = document.createRange()
  const lines: Array<{ top: number; height: number }> = []

  const walker = document.createTreeWalker(dom, NodeFilter.SHOW_TEXT)
  while (walker.nextNode()) {
    const text = walker.currentNode as Text
    if (!text.nodeValue?.length) continue
    range.selectNodeContents(text)
    for (const rect of range.getClientRects()) {
      if (rect.width <= 0 || rect.height <= 0) continue
      const relTop = rect.top - blockTop
      const last = lines[lines.length - 1]
      if (!last || Math.abs(last.top - relTop) > 2) {
        lines.push({ top: relTop, height: rect.height })
      } else {
        last.height = Math.max(last.height, rect.height)
      }
    }
  }

  if (lines.length === 0) {
    const h = dom.offsetHeight
    if (h > 0) lines.push({ top: 0, height: h })
  }

  return lines
}

function findPosAtVisualLine(
  view: EditorView,
  dom: HTMLElement,
  blockRect: DOMRect,
  relTop: number
): number | null {
  const coords = view.posAtCoords({
    left: blockRect.left + Math.min(Math.max(1, blockRect.width / 2), blockRect.width - 1),
    top: blockRect.top + relTop + 1,
  })
  return coords?.pos ?? null
}

function estimatedLineUnits(
  view: EditorView,
  node: PMNode,
  offset: number,
  dom: HTMLElement,
  totalH: number
): FlowUnit[] {
  const style = getComputedStyle(dom)
  let lineHeight = parseFloat(style.lineHeight)
  if (!Number.isFinite(lineHeight) || style.lineHeight === "normal") {
    lineHeight = (parseFloat(style.fontSize) || 16) * 1.25
  }
  const lineCount = Math.max(1, Math.ceil(totalH / lineHeight))
  const textLen = Math.max(1, node.content.size)
  const units: FlowUnit[] = []

  for (let i = 0; i < lineCount; i++) {
    const h =
      i === lineCount - 1
        ? Math.max(1, totalH - lineHeight * (lineCount - 1))
        : lineHeight
    const pos =
      i === 0
        ? offset
        : offset + Math.min(textLen - 1, Math.floor((textLen * i) / lineCount))
    units.push({ height: h, offset: pos })
  }

  return units
}

function expandTallLeafUnits(
  view: EditorView,
  node: PMNode,
  offset: number,
  dom: HTMLElement,
  totalH: number,
  bodyAreaH: number
): FlowUnit[] {
  if (totalH <= bodyAreaH) {
    return [{ height: totalH, offset }]
  }

  const lineSegments = measureVisualLineHeights(dom)
  if (lineSegments.length <= 1) {
    return estimatedLineUnits(view, node, offset, dom, totalH)
  }

  const blockRect = dom.getBoundingClientRect()
  return lineSegments.map((segment, lineIndex) => ({
    height: segment.height,
    offset:
      lineIndex === 0
        ? offset
        : (findPosAtVisualLine(view, dom, blockRect, segment.top) ?? offset),
  }))
}

function expandLeafBlockUnits(
  view: EditorView,
  node: PMNode,
  offset: number,
  units: FlowUnit[],
  bodyAreaH: number
): void {
  const dom = blockDomAtOffset(view, offset)
  if (!dom) return
  const totalH = measureBlockFlowHeight(dom)
  if (totalH <= 0) return
  units.push(...expandTallLeafUnits(view, node, offset, dom, totalH, bodyAreaH))
}

function expandNodeUnits(
  view: EditorView,
  node: PMNode,
  offset: number,
  units: FlowUnit[],
  bodyAreaH: number
): void {
  if (!node.isBlock) return

  if (LIST_TYPES.has(node.type.name)) {
    let itemOffset = offset + 1
    for (let j = 0; j < node.childCount; j++) {
      const item = node.child(j)
      const liDom = nestedDomAt(view, itemOffset, "li")
      const h = liDom ? measureBlockFlowHeight(liDom) : 0
      if (h > 0) {
        units.push(...expandTallLeafUnits(view, item, itemOffset, liDom, h, bodyAreaH))
      }
      itemOffset += item.nodeSize
    }
    return
  }

  if (node.type.name === "table") {
    let rowOffset = offset + 1
    for (let j = 0; j < node.childCount; j++) {
      const row = node.child(j)
      if (row.type.name !== "tableRow") {
        rowOffset += row.nodeSize
        continue
      }
      const rowDom = nestedDomAt(view, rowOffset, "tr")
      const h = rowDom ? measureBlockFlowHeight(rowDom) : 0
      if (h > 0) {
        units.push(...expandTallLeafUnits(view, row, rowOffset, rowDom, h, bodyAreaH))
      }
      rowOffset += row.nodeSize
    }
    return
  }

  if (node.type.name === "blockquote") {
    let innerOffset = offset + 1
    for (let j = 0; j < node.childCount; j++) {
      const child = node.child(j)
      if (child.isBlock) {
        expandLeafBlockUnits(view, child, innerOffset, units, bodyAreaH)
      }
      innerOffset += child.nodeSize
    }
    return
  }

  expandLeafBlockUnits(view, node, offset, units, bodyAreaH)
}

function collectFlowUnits(view: EditorView, bodyAreaH: number): FlowUnit[] {
  const units: FlowUnit[] = []
  const doc = view.state.doc
  let offset = 0

  for (let i = 0; i < doc.childCount; i++) {
    const node = doc.child(i)
    expandNodeUnits(view, node, offset, units, bodyAreaH)
    offset += node.nodeSize
  }

  return units
}

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) {
      const nextBodyStart = nextBreakY + interPageSpacer

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

      if (simulatedY < nextBodyStart) {
        const pushPx = nextBodyStart - simulatedY
        pushes.push({ blockIndex, pushPx, breakY: nextBreakY })
        simulatedY = nextBodyStart
        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) {
      const nextBodyStart = nextBreakY + interPageSpacer

      if (simulatedY < nextBreakY) {
        simulatedY = nextBodyStart
        nextBreakY += pageStep
        break
      }

      if (simulatedY < nextBodyStart) {
        simulatedY = nextBodyStart
        nextBreakY += pageStep
        break
      }

      nextBreakY += pageStep
    }
    simulatedY += blockH
  }

  return simulatedY
}

/** 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 units = collectFlowUnits(view, bodyAreaH)

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

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

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

/** Measure paginated prose height using the same flow units as layout. */
export function measureFlowContentHeight(view: EditorView): number {
  const metrics = readPageFlowMetrics(view.dom)
  if (!metrics) return 0

  const units = collectFlowUnits(view, metrics.bodyAreaH)
  if (units.length === 0) return 0

  const simulated = computeSimulatedLayoutHeight(
    units.map((u) => ({ height: u.height })),
    metrics.bodyAreaH,
    metrics.interPageSpacer
  )

  let maxBottom = 0
  for (const child of view.dom.children) {
    const el = child as HTMLElement
    maxBottom = Math.max(maxBottom, el.offsetTop + el.offsetHeight)
  }

  return Math.max(simulated, maxBottom)
}

/** 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
          },
        },
      }),
    ]
  },
})