import * as React from 'react'
import { cn } from '../lib/cn'

function countRowColumns(children: React.ReactNode): number {
  let count = 0
  React.Children.forEach(children, (child) => {
    if (React.isValidElement(child)) {
      const span = (child.props as { colSpan?: number }).colSpan
      count += typeof span === 'number' && span > 0 ? span : 1
    }
  })
  return count
}

interface TableGridContextValue {
  columnCount: number
  registerColumnCount: (count: number) => void
}

const TableGridContext = React.createContext<TableGridContextValue | null>(null)

function useTableGrid(): TableGridContextValue {
  const ctx = React.useContext(TableGridContext)
  if (!ctx) {
    throw new Error('Table row/cell components must be used within Table')
  }
  return ctx
}

export const Table = React.forwardRef<
  HTMLTableElement,
  React.TableHTMLAttributes<HTMLTableElement>
>(({ className, ...props }, ref) => {
  const columnCountRef = React.useRef(0)
  const [, bump] = React.useReducer((n: number) => n + 1, 0)

  const registerColumnCount = React.useCallback((count: number) => {
    if (count > columnCountRef.current) {
      columnCountRef.current = count
      bump()
    }
  }, [])

  const columnCount = columnCountRef.current

  const contextValue = React.useMemo(
    () => ({ columnCount, registerColumnCount }),
    [columnCount, registerColumnCount],
  )

  return (
    <TableGridContext.Provider value={contextValue}>
      <div className="w-full overflow-auto" tabIndex={0}>
        <div
          ref={ref as React.Ref<HTMLDivElement>}
          role="table"
          className={cn('w-full text-body-2 text-ink', className)}
          {...props}
        />
      </div>
    </TableGridContext.Provider>
  )
})
Table.displayName = 'Table'

export const Thead = React.forwardRef<
  HTMLTableSectionElement,
  React.HTMLAttributes<HTMLTableSectionElement>
>(({ className, ...props }, ref) => (
  <div
    ref={ref as React.Ref<HTMLDivElement>}
    role="rowgroup"
    className={cn('border-b border-line', className)}
    {...props}
  />
))
Thead.displayName = 'Thead'

export const Tbody = React.forwardRef<
  HTMLTableSectionElement,
  React.HTMLAttributes<HTMLTableSectionElement>
>(({ className, ...props }, ref) => (
  <div
    ref={ref as React.Ref<HTMLDivElement>}
    role="rowgroup"
    className={cn(className)}
    {...props}
  />
))
Tbody.displayName = 'Tbody'

export const Tr = React.forwardRef<
  HTMLTableRowElement,
  React.HTMLAttributes<HTMLTableRowElement>
>(({ className, children, style, ...props }, ref) => {
  const { columnCount, registerColumnCount } = useTableGrid()
  const rowColumnCount = countRowColumns(children)

  React.useLayoutEffect(() => {
    if (rowColumnCount > 0) {
      registerColumnCount(rowColumnCount)
    }
  }, [rowColumnCount, registerColumnCount])

  const gridTemplateColumns =
    columnCount > 0 ? `repeat(${columnCount}, minmax(0, 1fr))` : undefined

  return (
    <div
      ref={ref as React.Ref<HTMLDivElement>}
      role="row"
      className={cn('border-b border-line-subtle', className)}
      style={
        gridTemplateColumns
          ? { display: 'grid', gridTemplateColumns, ...style }
          : style
      }
      {...props}
    >
      {children}
    </div>
  )
})
Tr.displayName = 'Tr'

export const Th = React.forwardRef<
  HTMLTableCellElement,
  React.ThHTMLAttributes<HTMLTableCellElement>
>(({ className, scope: _scope, colSpan, rowSpan: _rowSpan, style, ...props }, ref) => {
  const span = colSpan ?? 1
  const cellStyle: React.CSSProperties = {
    ...style,
    ...(span > 1 ? { gridColumn: `span ${span}` } : {}),
  }

  return (
    <div
      ref={ref as React.Ref<HTMLDivElement>}
      role="columnheader"
      className={cn('min-h-8 min-w-0 px-2 py-2 text-start text-meta font-medium text-ink-soft [overflow-wrap:anywhere]', className)}
      aria-colspan={span > 1 ? span : undefined}
      style={cellStyle}
      {...props}
    />
  )
})
Th.displayName = 'Th'

export const Td = React.forwardRef<
  HTMLTableCellElement,
  React.TdHTMLAttributes<HTMLTableCellElement>
>(({ className, colSpan, style, ...props }, ref) => {
  const span = colSpan ?? 1
  const cellStyle: React.CSSProperties = {
    ...style,
    ...(span > 1 ? { gridColumn: `span ${span}` } : {}),
  }

  return (
    <div
      ref={ref as React.Ref<HTMLDivElement>}
      role="cell"
      className={cn('min-h-8 min-w-0 px-2 py-2 [overflow-wrap:anywhere]', className)}
      aria-colspan={span > 1 ? span : undefined}
      style={cellStyle}
      {...props}
    />
  )
})
Td.displayName = 'Td'