Spaces:

MohamedSameh77i
/

Agent_PDF

Sleeping

File size: 10,109 Bytes

import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
import { useStudio } from './StudioContext'
import { getTablesForPage } from './studioUtils'
import TableEditor from './TableEditorFull'

function toNumber(value, fallback = 0) {
  const num = Number(value)
  return Number.isFinite(num) ? num : fallback
}

function makeCellId(tableId, row, col) {
  return `t${tableId}-r${row}-c${col}`
}

function buildGridFromTable(table) {
  const cells = table?.cells || []
  // Safely compute dimensions — guard against undefined/null span fields.
  const rowCount = cells.length
    ? Math.max(...cells.map((cell) => toNumber(cell.row, 0) + Math.max(1, toNumber(cell.row_span, 1))))
    : 1
  const colCount = cells.length
    ? Math.max(...cells.map((cell) => toNumber(cell.col, 0) + Math.max(1, toNumber(cell.col_span, 1))))
    : 1

  const safRows = Math.max(1, rowCount)
  const safCols = Math.max(1, colCount)

  const grid = Array.from({ length: safRows }, (_, row) =>
    Array.from({ length: safCols }, (_, col) => ({
      id: makeCellId(table?.table_id ?? 0, row, col),
      value: '',
      rowSpan: 1,
      colSpan: 1,
      hidden: false,
    }))
  )

  for (const cell of cells) {
    const row = toNumber(cell.row, 0)
    const col = toNumber(cell.col, 0)
    const rowSpan = Math.max(1, toNumber(cell.row_span, 1))
    const colSpan = Math.max(1, toNumber(cell.col_span, 1))
    // Skip cells that are out-of-bounds (shouldn't happen but guards against bad data).
    if (row >= safRows || col >= safCols || !grid[row]?.[col]) {
      continue
    }
    grid[row][col] = {
      id: makeCellId(table.table_id, row, col),
      value: String(cell.text || ''),
      rowSpan,
      colSpan,
      hidden: false,
      fontFamily: String(cell.font_family || ''),
      fontWeight: String(cell.font_weight || ''),
      backgroundColor: String(cell.background_class || ''),
    }
    // Mark spanned-over slots as hidden
    for (let r = row; r < Math.min(row + rowSpan, safRows); r += 1) {
      for (let c = col; c < Math.min(col + colSpan, safCols); c += 1) {
        if ((r !== row || c !== col) && grid[r]?.[c]) {
          grid[r][c].hidden = true
          grid[r][c].value = ''
          grid[r][c].rowSpan = 1
          grid[r][c].colSpan = 1
        }
      }
    }
  }

  return grid
}

function normalizeCellsFromGrid(grid, table) {
  if (!Array.isArray(grid) || !grid.length || !Array.isArray(grid[0]) || !grid[0].length) {
    return []
  }

  const previousByAnchor = new Map(
    (table?.cells || []).map((cell) => [`${cell.row},${cell.col}`, cell]),
  )
  const rowCount = grid.length
  const colCount = grid[0].length
  const [tx1, ty1, tx2, ty2] = Array.isArray(table?.bbox) && table.bbox.length === 4
    ? table.bbox.map((value) => toNumber(value, 0))
    : [0, 0, colCount, rowCount]
  const cellWidth = (tx2 - tx1) / Math.max(colCount, 1)
  const cellHeight = (ty2 - ty1) / Math.max(rowCount, 1)

  const nextCells = []
  for (let row = 0; row < rowCount; row += 1) {
    for (let col = 0; col < colCount; col += 1) {
      const current = grid[row][col]
      if (!current || current.hidden) {
        continue
      }
      const rowSpan = Math.max(1, toNumber(current.rowSpan, 1))
      const colSpan = Math.max(1, toNumber(current.colSpan, 1))
      const previous = previousByAnchor.get(`${row},${col}`)
      const fallbackBbox = [
        Number((tx1 + (col * cellWidth)).toFixed(1)),
        Number((ty1 + (row * cellHeight)).toFixed(1)),
        Number((tx1 + ((col + colSpan) * cellWidth)).toFixed(1)),
        Number((ty1 + ((row + rowSpan) * cellHeight)).toFixed(1)),
      ]

      nextCells.push({
        row,
        col,
        row_span: rowSpan,
        col_span: colSpan,
        text: String(current.value || ''),
        bbox: Array.isArray(previous?.bbox) && previous.bbox.length === 4 ? previous.bbox : fallbackBbox,
        ocr_score: previous?.ocr_score ?? null,
        font_family: String(current.fontFamily || previous?.font_family || ''),
        font_weight: String(current.fontWeight || previous?.font_weight || ''),
        background_class: String(current.backgroundColor || previous?.background_class || ''),
      })
    }
  }

  return nextCells
}

function hashGrid(grid) {
  return JSON.stringify(
    (grid || []).map((row) => row.map((cell) => ({
      v: cell.value,
      rs: cell.rowSpan,
      cs: cell.colSpan,
      h: cell.hidden,
      bg: cell.backgroundColor || '',
      fw: cell.fontWeight || '',
      ff: cell.fontFamily || '',
    }))),
  )
}

export default function StudioGridPane({
  annotation,
  onReplaceTable,
  onEditingChange,
  hoveredCell,
  onHoverCellChange,
}) {
  const { state, dispatch } = useStudio()
  const containerRef = useRef(null)
  const saveTimerRef = useRef(null)
  const activeHashRef = useRef('')
  const latestGridRef = useRef([])
  const contextTableRef = useRef(null)
  const isSavingRef = useRef(false)
  const [isSaving, setIsSaving] = useState(false)

  const activePageTables = useMemo(
    () => getTablesForPage(annotation, state.activePage),
    [annotation, state.activePage],
  )
  const contextTable = activePageTables.find((table) => table.table_id === state.activeTable) || activePageTables[0] || null
  const hoveredCellInContext = hoveredCell && contextTable && hoveredCell.tableId === contextTable.table_id
    ? { r: hoveredCell.row, c: hoveredCell.col }
    : null
  const editorSeed = useMemo(
    () => (contextTable ? buildGridFromTable(contextTable) : []),
    [contextTable],
  )

  useEffect(() => {
    if (contextTable && state.activeTable !== contextTable.table_id) {
      dispatch({ type: 'setActiveTable', tableId: contextTable.table_id })
    }
  }, [contextTable, dispatch, state.activeTable])

  // Keep a ref to contextTable so persistGrid (debounced) always uses the latest value.
  useEffect(() => {
    contextTableRef.current = contextTable
  }, [contextTable])

  useEffect(() => {
    const nextHash = hashGrid(editorSeed)
    activeHashRef.current = nextHash
    latestGridRef.current = editorSeed
  }, [editorSeed])

  useEffect(() => {
    const node = containerRef.current
    if (!node) {
      return undefined
    }
    const handleFocusIn = () => onEditingChange?.(true)
    const handleFocusOut = () => {
      window.requestAnimationFrame(() => {
        const active = document.activeElement
        onEditingChange?.(Boolean(active && node.contains(active)))
      })
    }
    node.addEventListener('focusin', handleFocusIn)
    node.addEventListener('focusout', handleFocusOut)
    return () => {
      node.removeEventListener('focusin', handleFocusIn)
      node.removeEventListener('focusout', handleFocusOut)
      onEditingChange?.(false)
    }
  }, [onEditingChange])

  useEffect(() => () => {
    if (saveTimerRef.current) {
      window.clearTimeout(saveTimerRef.current)
    }
  }, [])

  const persistGrid = useCallback(async () => {
    const table = contextTableRef.current
    if (!table || !onReplaceTable || isSavingRef.current) {
      return
    }
    const latest = latestGridRef.current
    const nextHash = hashGrid(latest)
    if (nextHash === activeHashRef.current) {
      return
    }
    isSavingRef.current = true
    setIsSaving(true)
    try {
      const nextCells = normalizeCellsFromGrid(latest, table)
      const ok = await onReplaceTable({
        tableId: table.table_id,
        cells: nextCells,
      })
      if (ok) {
        activeHashRef.current = nextHash
      }
    } finally {
      isSavingRef.current = false
      setIsSaving(false)
    }
  }, [onReplaceTable])

  const handleGridChange = useCallback((nextGrid) => {
    latestGridRef.current = nextGrid
    if (saveTimerRef.current) {
      window.clearTimeout(saveTimerRef.current)
    }
    saveTimerRef.current = window.setTimeout(() => {
      void persistGrid()
    }, 350)
  }, [persistGrid])

  if (!activePageTables.length) {
    return (
      <section className="studio-panel studio-grid-pane">
        <div className="studio-grid-context">
          <span>Viewing: Page {state.activePage + 1}</span>
          <span>No table on this page</span>
        </div>
        <div className="studio-grid-scroll">
          <div className="studio-empty-state">
            No tables were detected on page {state.activePage + 1}.
          </div>
        </div>
      </section>
    )
  }

  return (
    <section className="studio-panel studio-grid-pane" ref={containerRef}>
      <div className="studio-grid-context">
        <span>Viewing: Page {state.activePage + 1}</span>
        <div style={{ display: 'flex', alignItems: 'center', gap: 12 }}>
          <label htmlFor="studio-table-select">Table</label>
          <select
            id="studio-table-select"
            value={contextTable?.table_id ?? ''}
            onChange={(event) => {
              dispatch({ type: 'setActiveTable', tableId: Number(event.target.value) })
            }}
          >
            {activePageTables.map((table) => (
              <option key={table.table_id} value={table.table_id}>
                Table {table.table_id + 1}
              </option>
            ))}
          </select>
          {isSaving && <span>Saving…</span>}
        </div>
      </div>

      <div className="studio-grid-scroll">
        {editorSeed.length ? (
          <TableEditor
            initialGrid={editorSeed}
            onGridChange={handleGridChange}
            persistKey={`${state.activePage}:${contextTable?.table_id ?? 'none'}`}
            hoveredCell={hoveredCellInContext}
            onHoverCellChange={(cell) => {
              if (!onHoverCellChange || !contextTable) {
                return
              }
              onHoverCellChange(cell ? {
                tableId: contextTable.table_id,
                row: cell.r,
                col: cell.c,
                pageIndex: contextTable.page ?? state.activePage,
              } : null)
            }}
          />
        ) : (
          <div className="studio-empty-state">Selected table has no cells.</div>
        )}
      </div>
    </section>
  )
}