"use client"
import {
  useVirtualizer,
  useWindowVirtualizer,
  VirtualizerOptions,
} from "@tanstack/react-virtual"
import React, { useCallback, useRef } from "react"
import { useCallbackRef, useDedupedItems } from "../../hooks"
import type { Awaitable } from "../../types/utils"
import { TableProps } from "../Table"
import { TABLE_SIZES } from "../Table/constants"
import { TableHeader } from "../Table/TableHeader"
import { TableProvider } from "../Table/TableProvider"
import { TableRow } from "../Table/TableRow"
import { VirtualizedTableView } from "./VirtualizedTableView"

type VirtualizerProps = Pick<
  VirtualizerOptions<Window, Element>,
  "initialRect"
> & {
  overscanBy?: number
}

type VirtualizedTablePaginationProps = {
  hasNext: boolean
  loadNext: () => Awaitable<unknown>
  /**
   * The threshold at which to pre-fetch data. A threshold X means that new data
   * should start loading when a user scrolls within X cells of the end of the
   * `items` array.
   *
   * @defaultValue `DEFAULT_PAGINATION_THRESHOLD`
   */
  paginationThreshold?: number
  paginationError?: boolean
}

type VirtualizedTableBaseProps<T, C> = {
  /**
   * Determines whether to use window or direct table scrolling. In `window`
   * mode, scrolling the browser window will cause the table to scroll even when
   * the user is not hovering over the table html element directly. In `table`
   * mode, only scrolling on top of the table will cause the table to scroll.
   *
   *
   * @defaultValue `"window"`
   */
  scrollMode?: "window" | "table"
  /**
   * Ref to the container that should be scrolled. This is required when using table mode if a height is not set.
   */
  scrollContainerRef?: React.RefObject<HTMLDivElement | null>
  tableBodyRef?: React.RefObject<HTMLDivElement | null>
  setTableBodyRef?: (ref: HTMLDivElement | null) => void
  id?: string

  /**
   * Override the default size of the table.
   */
  overrides?: {
    size?: {
      gap?: number
      height?: number
    }
  }
} & Pick<
  TableProps<T, C>,
  | "items"
  | "className"
  | "contentClassName"
  | "style"
  | "contentStyle"
  | "itemKey"
  | "header"
  | "renderRow"
  | "context"
  | "size"
  | "headerBorderOnScroll"
  | "dividers"
> &
  VirtualizerProps &
  Partial<VirtualizedTablePaginationProps>

export type VirtualizedTableModeProps<T, C> = VirtualizedTableBaseProps<
  T,
  C
> & {
  // For scroll mode, you must either set a container ref or the set a height on the table.
  scrollMode: "table"
}

export type VirtualizedTableWindowModeProps<
  T,
  C = unknown,
> = VirtualizedTableBaseProps<T, C> & {
  scrollMode?: "window"
}

export type VirtualizedTableProps<T, C = unknown> =
  | VirtualizedTableModeProps<T, C>
  | VirtualizedTableWindowModeProps<T, C>

const isTableMode = <T, C>(
  props: Partial<VirtualizedTableProps<T, C>>,
): props is VirtualizedTableModeProps<T, C> => {
  return props.scrollMode === "table"
}

const VirtualizedTableBase = <T, C = unknown>({
  items,
  ...props
}: VirtualizedTableProps<T, C>) => {
  const itemKeyRef = useRef(props.itemKey)
  const dedupedItems = useDedupedItems(items, props.itemKey)

  if (isTableMode(props)) {
    return (
      <TableScrollVirtualizedTable
        {...props}
        itemKey={itemKeyRef.current}
        items={dedupedItems}
      />
    )
  }

  return (
    <WindowScrollVirtualizedTable
      {...props}
      itemKey={itemKeyRef.current}
      items={dedupedItems}
    />
  )
}

const getDefaultOverscan = (rowHeight: number): number => {
  // 6 rows at 32px height
  return Math.ceil(190 / rowHeight)
}

const WindowScrollVirtualizedTable = <T, C>({
  items,
  overscanBy,
  tableBodyRef,
  setTableBodyRef: setTableBodyRefProp,
  initialRect,
  itemKey,
  size = "md",
  overrides,
  dividers,
  ...props
}: VirtualizedTableProps<T, C>) => {
  const [tableRef, setTableRefCallback] =
    useCallbackRef<HTMLDivElement>(tableBodyRef)

  const getItemKey = useCallback(
    (index: number) => itemKey?.(items[index], index) ?? index,
    [itemKey, items],
  )

  const tableSize = TABLE_SIZES[size]
  const gap = overrides?.size?.gap ?? tableSize.gap
  const gapAccountingForDividers = dividers ? 0 : gap

  const rowHeight = overrides?.size?.height ?? tableSize.rowHeight

  const overscan = overscanBy ?? getDefaultOverscan(rowHeight)

  const rowVirtualizer = useWindowVirtualizer({
    count: items.length,
    estimateSize: () => rowHeight,
    getItemKey,
    overscan,
    scrollMargin: tableRef.current?.offsetTop,
    initialRect,
    gap: gapAccountingForDividers,
  })

  const setTableRef = (ref: HTMLDivElement) => {
    if (!ref || ref === tableRef.current) {
      return
    }

    setTableRefCallback(ref)
    setTableBodyRefProp?.(ref)
  }

  return (
    <TableProvider dividers={dividers} size={size}>
      <VirtualizedTableView
        gap={gapAccountingForDividers}
        itemKey={itemKey}
        items={items}
        rowVirtualizer={rowVirtualizer}
        setTableRef={setTableRef}
        tableRef={tableRef}
        {...props}
      />
    </TableProvider>
  )
}

const TableScrollVirtualizedTable = <T, C>({
  items,
  overscanBy,
  scrollContainerRef,
  tableBodyRef,
  setTableBodyRef: setTableBodyRefProp,
  initialRect,
  className,
  size = "md",
  overrides,
  itemKey,
  dividers,
  ...props
}: VirtualizedTableModeProps<T, C>) => {
  const [tableRef, setTableRefCallback] =
    useCallbackRef<HTMLDivElement>(tableBodyRef)

  const tableSize = TABLE_SIZES[size]
  const gap = overrides?.size?.gap ?? tableSize.gap
  const gapAccountingForDividers = dividers ? 0 : gap
  const rowHeight = overrides?.size?.height ?? tableSize.rowHeight

  const overscan = overscanBy ?? getDefaultOverscan(rowHeight)

  const getItemKey = useCallback(
    (index: number) => itemKey?.(items[index], index) ?? index,
    [itemKey, items],
  )

  const rowVirtualizer = useVirtualizer({
    count: items.length,
    estimateSize: () => rowHeight,
    getItemKey,
    getScrollElement: () => scrollContainerRef?.current ?? tableRef.current,
    overscan,
    initialRect,
    scrollMargin: tableRef.current?.offsetTop,
    gap: gapAccountingForDividers,
  })

  const setTableRef = (ref: HTMLDivElement) => {
    if (!ref || ref === tableRef.current) {
      return
    }

    setTableRefCallback(ref)
    setTableBodyRefProp?.(ref)
  }

  return (
    <TableProvider dividers={dividers} size={size}>
      <VirtualizedTableView
        className={className}
        gap={gapAccountingForDividers}
        itemKey={itemKey}
        items={items}
        rowVirtualizer={rowVirtualizer}
        setTableRef={setTableRef}
        tableRef={tableRef}
        {...props}
      />
    </TableProvider>
  )
}

/**
 * A table that automatically virtualizes its rows to improve performance.
 * Rows that are out of the viewport are not rendered.
 */
export const VirtualizedTable = Object.assign(VirtualizedTableBase, {
  Header: TableHeader,
  Row: TableRow,
})