<script setup lang="ts" generic="TData, TValue">
import {
  FlexRender,
  getExpandedRowModel,
  getCoreRowModel,
  getFilteredRowModel,
  getPaginationRowModel,
  getSortedRowModel,
  useVueTable,
  type ColumnOrderState,
  type ColumnFiltersState,
  type ExpandedState,
  type Header,
  type Row,
  type SortingState,
  type Table,
  type Updater,
} from '@tanstack/vue-table';
import {
  UiTable,
  UiTableBody,
  UiTableCell,
  UiTableHead,
  UiTableHeader,
  UiTableRow,
} from '@/components/UiTable';
import { valueUpdater } from '@/components/core/table/utils';
import {
  heightTransitionClass,
  onHeightTransitionAfterEnter,
  onHeightTransitionBeforeLeave,
  onHeightTransitionEnter,
  onHeightTransitionLeave,
} from '@/lib/heightTransition';
import { ref, nextTick, onMounted, onUnmounted } from 'vue';
import type { UiDataTableProps } from './types';
import { DATA_TABLE_LABELS, DEFAULT_COLUMN_SIZE } from './constants';
import { useDrag } from '@/lib/useDrag';

const INTERACTIVE_SELECTOR =
  'button, a, input, textarea, select, [role="button"], [role="link"], [role="menuitem"], [role="checkbox"], [role="switch"], [contenteditable="true"]';
const DRAGGED_HEADER_CLASS = 'border-2 border-border-surface-focus-ring';
const DRAGGED_COLUMN_CELL_CLASS = 'border-x-2 border-border-surface-focus-ring';
const DRAG_GHOST_CLASS =
  'rounded-sm border-2 border-border-surface-focus-ring bg-background-accent-soft/10';
const DROP_LINE_WIDTH_PX = 2;
type DropSide = 'left' | 'right';

type UiDataTableExpandedRowSlotProps<TData> = {
  /** The expanded TanStack row instance. */
  row: Row<TData>;
  /** The TanStack table instance. */
  table: Table<TData>;
};

defineOptions({
  name: 'UiDataTable',
});

const props = withDefaults(defineProps<UiDataTableProps<TData, TValue>>(), {
  emptyText: DATA_TABLE_LABELS.emptyText,
  draggableColumns: false,
});

defineSlots<{
  /** Content rendered in the empty state cell. */
  empty(): unknown;
  /** Content rendered in a separate table row after an expanded data row. */
  'expanded-row'(props: UiDataTableExpandedRowSlotProps<TData>): unknown;
}>();

const sorting = ref<SortingState>([]);
const columnOrder = ref<ColumnOrderState>([]);
const columnFilters = ref<ColumnFiltersState>([]);
const expanded = ref<ExpandedState>({});
const expandedRowsToRender = ref<Record<string, true>>({});
const rowSelection = ref({});
const globalFilter = ref('');
const draggingColumnId = ref<string>();
const wrapperRef = ref<HTMLElement>();
const focusedColumnIndicator = ref<{
  /** Horizontal offset from the table wrapper. */
  left: number;
  /** Vertical offset from the table wrapper. */
  top: number;
  /** Width of the focused column. */
  width: number;
  /** Height of the rendered table. */
  height: number;
}>();
const dropIndicator = ref<{
  /** Horizontal offset from the table wrapper. */
  left: number;
  /** Vertical offset from the table wrapper. */
  top: number;
  /** Height of the rendered table. */
  height: number;
  /** Edge used as insertion point for the current drop target. */
  side: DropSide;
}>();
const activeDropTarget = ref<{
  /** Column currently receiving the dragged column. */
  columnId: string;
  /** Edge used as insertion point for the active drop target. */
  side: DropSide;
}>();
const { cleanupDragGhost, getDragData, setDragData, setDragGhost } = useDrag();
let isPointerFocusInteraction = false;
let resetPointerFocusInteractionTimeout: number | undefined;

function handleColumnOrderChange(updaterOrValue: Updater<ColumnOrderState>) {
  valueUpdater(updaterOrValue, columnOrder);
  props.onColumnOrderChange?.(columnOrder.value);
}

function syncExpandedRowsToRender() {
  const nextRowsToRender = { ...expandedRowsToRender.value };

  for (const row of table.getRowModel().rows) {
    if (row.getIsExpanded()) nextRowsToRender[row.id] = true;
  }

  expandedRowsToRender.value = nextRowsToRender;
}

function handleExpandedChange(updaterOrValue: Updater<ExpandedState>) {
  valueUpdater(updaterOrValue, expanded);
  nextTick(syncExpandedRowsToRender);
}

const table = useVueTable({
  get data() {
    return props.data;
  },
  get columns() {
    return props.columns;
  },
  getRowId: props.getRowId,
  getCoreRowModel: getCoreRowModel(),
  getPaginationRowModel: getPaginationRowModel(),
  getSortedRowModel: getSortedRowModel(),
  getFilteredRowModel: getFilteredRowModel(),
  getExpandedRowModel: getExpandedRowModel(),
  getRowCanExpand: props.getRowCanExpand,
  enableRowSelection: true,
  onSortingChange: (updaterOrValue) => valueUpdater(updaterOrValue, sorting),
  onColumnOrderChange: handleColumnOrderChange,
  onColumnFiltersChange: (updaterOrValue) =>
    valueUpdater(updaterOrValue, columnFilters),
  onExpandedChange: handleExpandedChange,
  onRowSelectionChange: (updaterOrValue) =>
    valueUpdater(updaterOrValue, rowSelection),
  onGlobalFilterChange: (updaterOrValue) =>
    valueUpdater(updaterOrValue, globalFilter),
  state: {
    get sorting() {
      return sorting.value;
    },
    get columnOrder() {
      return columnOrder.value;
    },
    get columnFilters() {
      return columnFilters.value;
    },
    get expanded() {
      return expanded.value;
    },
    get rowSelection() {
      return rowSelection.value;
    },
    get globalFilter() {
      return globalFilter.value;
    },
  },
});

function canDragHeader(header: Header<TData, unknown>) {
  return (
    props.draggableColumns && !header.isPlaceholder && header.colSpan === 1
  );
}

function getColumnHeaderDragStateClass(columnId: string) {
  if (!props.draggableColumns) return undefined;

  return [
    draggingColumnId.value === columnId ? DRAGGED_HEADER_CLASS : undefined,
  ];
}

function getColumnCellDragStateClass(columnId: string, isLastRow: boolean) {
  if (!props.draggableColumns) return undefined;

  return [
    draggingColumnId.value === columnId
      ? [
          DRAGGED_COLUMN_CELL_CLASS,
          isLastRow ? 'border-b-2 border-border-surface-focus-ring' : undefined,
        ]
      : undefined,
  ];
}

function getColumnHeaderInteractionClass(columnId: string) {
  if (!props.draggableColumns) return undefined;

  if (draggingColumnId.value === columnId) {
    return 'cursor-grabbing';
  }

  return 'cursor-grab';
}

function getDropSide(
  target: HTMLElement,
  event: DragEvent,
  fallbackSide: DropSide = 'right',
) {
  if (!Number.isFinite(event.clientX)) return fallbackSide;

  const targetRect = target.getBoundingClientRect();
  const midpointX = targetRect.left + targetRect.width / 2;
  return event.clientX > midpointX ? 'right' : 'left';
}

function updateDropIndicator(target: HTMLElement, side: DropSide = 'left') {
  const indicator = getDropIndicator(target, side);
  if (!indicator) return;
  dropIndicator.value = indicator;
}

function getColumnIndicator(target: HTMLElement) {
  const wrapper = wrapperRef.value;
  const tableElement = wrapper?.querySelector<HTMLElement>(
    '[data-slot="table"]',
  );
  if (!wrapper || !tableElement) return undefined;

  const targetRect = target.getBoundingClientRect();
  const wrapperRect = wrapper.getBoundingClientRect();
  const tableRect = tableElement.getBoundingClientRect();

  return {
    left: targetRect.left - wrapperRect.left,
    top: tableRect.top - wrapperRect.top,
    width: targetRect.width,
    height: tableRect.height,
  };
}

function getDropIndicator(target: HTMLElement, side: DropSide) {
  const columnIndicator = getColumnIndicator(target);
  if (!columnIndicator) return undefined;

  const edgeX =
    side === 'right'
      ? columnIndicator.left + columnIndicator.width
      : columnIndicator.left;

  return {
    left: edgeX - DROP_LINE_WIDTH_PX / 2,
    top: columnIndicator.top,
    height: columnIndicator.height,
    side,
  };
}

function updateFocusedColumnIndicator(target: HTMLElement) {
  focusedColumnIndicator.value = getColumnIndicator(target);
}

function handleColumnHeaderFocus(event: FocusEvent) {
  if (isPointerFocusInteraction) return;

  updateFocusedColumnIndicator(event.currentTarget as HTMLElement);
}

function handleColumnHeaderClick() {
  focusedColumnIndicator.value = undefined;
}

function handleColumnHeaderBlur() {
  focusedColumnIndicator.value = undefined;
}

function clearFocusedColumnIndicatorForPointerAction() {
  isPointerFocusInteraction = true;
  focusedColumnIndicator.value = undefined;

  if (resetPointerFocusInteractionTimeout !== undefined) {
    window.clearTimeout(resetPointerFocusInteractionTimeout);
  }

  resetPointerFocusInteractionTimeout = window.setTimeout(() => {
    isPointerFocusInteraction = false;
    resetPointerFocusInteractionTimeout = undefined;
  }, 0);
}

function clearResetPointerFocusInteractionTimeout() {
  if (resetPointerFocusInteractionTimeout === undefined) return;

  window.clearTimeout(resetPointerFocusInteractionTimeout);
  resetPointerFocusInteractionTimeout = undefined;
}

function getColumnIds() {
  return table.getAllLeafColumns().map((column) => column.id);
}

function getColumnHeaderElement(columnId: string) {
  const columnHeaders = wrapperRef.value?.querySelectorAll<HTMLElement>(
    'thead [data-column-id]',
  );

  return Array.from(columnHeaders ?? []).find(
    (columnHeader) => columnHeader.dataset.columnId === columnId,
  );
}

function reorderColumn(
  draggedColumnId: string,
  targetColumnId: string,
  side: DropSide = 'left',
) {
  const columnIds = getColumnIds();
  const draggedIndex = columnIds.indexOf(draggedColumnId);
  const targetIndex = columnIds.indexOf(targetColumnId);
  if (draggedIndex === -1 || targetIndex === -1) return false;

  let insertionIndex = side === 'right' ? targetIndex + 1 : targetIndex;
  if (draggedIndex < insertionIndex) {
    insertionIndex -= 1;
  }
  if (insertionIndex === draggedIndex) return false;

  const nextColumnOrder = [...columnIds];
  const [draggedColumn] = nextColumnOrder.splice(draggedIndex, 1);
  if (draggedColumn === undefined) return false;
  nextColumnOrder.splice(insertionIndex, 0, draggedColumn);

  if (nextColumnOrder === columnIds) return false;

  table.setColumnOrder(nextColumnOrder);
  return true;
}

function finishColumnDrop() {
  draggingColumnId.value = undefined;
  activeDropTarget.value = undefined;
  focusedColumnIndicator.value = undefined;
  dropIndicator.value = undefined;
  cleanupDragGhost();
}

onMounted(() => {
  document.addEventListener(
    'pointerdown',
    clearFocusedColumnIndicatorForPointerAction,
    true,
  );
});

onUnmounted(() => {
  document.removeEventListener(
    'pointerdown',
    clearFocusedColumnIndicatorForPointerAction,
    true,
  );
  clearResetPointerFocusInteractionTimeout();
  cleanupDragGhost();
  focusedColumnIndicator.value = undefined;
});

function clampDragOffset(offset: number, size: number) {
  if (!Number.isFinite(offset)) return size / 2;

  return Math.max(0, Math.min(offset, size));
}

function getDragGhostOptions(event: DragEvent, target: HTMLElement) {
  const tableElement = wrapperRef.value?.querySelector<HTMLElement>(
    '[data-slot="table"]',
  );
  const targetRect = target.getBoundingClientRect();
  const tableRect = tableElement?.getBoundingClientRect();
  const height = tableRect?.height ?? targetRect.height;

  return {
    width: targetRect.width,
    height,
    offsetX: clampDragOffset(event.clientX - targetRect.left, targetRect.width),
    offsetY: clampDragOffset(
      event.clientY - (tableRect?.top ?? targetRect.top),
      height,
    ),
  };
}

function startColumnDrag(
  event: DragEvent,
  columnId: string,
  target: HTMLElement,
) {
  focusedColumnIndicator.value = undefined;
  draggingColumnId.value = columnId;
  activeDropTarget.value = undefined;
  setDragData(event, columnId);

  if (event.dataTransfer) event.dataTransfer.effectAllowed = 'move';

  setDragGhost(event, {
    ...getDragGhostOptions(event, target),
    className: DRAG_GHOST_CLASS,
  });
  nextTick(() => updateDropIndicator(target, getDropSide(target, event)));
}

function handleColumnDragStart(
  event: DragEvent,
  header: Header<TData, unknown>,
) {
  if (!canDragHeader(header)) return;

  startColumnDrag(event, header.column.id, event.currentTarget as HTMLElement);
}

function handleColumnCellDragStart(event: DragEvent, columnId: string) {
  if (!props.draggableColumns) return;

  startColumnDrag(event, columnId, event.currentTarget as HTMLElement);
}

function handleColumnKeydown(
  event: KeyboardEvent,
  header: Header<TData, unknown>,
) {
  if (!canDragHeader(header)) return;

  if (event.key !== 'ArrowLeft' && event.key !== 'ArrowRight') return;

  const direction = event.key === 'ArrowLeft' ? -1 : 1;

  const columnIds = getColumnIds();
  const currentIndex = columnIds.indexOf(header.column.id);
  const targetColumnId = columnIds[currentIndex + direction];
  const targetColumnHeader = targetColumnId
    ? getColumnHeaderElement(targetColumnId)
    : undefined;

  if (!targetColumnId || !targetColumnHeader) return;

  event.preventDefault();
  draggingColumnId.value = header.column.id;
  updateDropIndicator(targetColumnHeader, direction === 1 ? 'right' : 'left');

  if (
    reorderColumn(
      header.column.id,
      targetColumnId,
      direction === 1 ? 'right' : 'left',
    )
  ) {
    nextTick(() => {
      const movedColumnHeader = getColumnHeaderElement(header.column.id);
      if (!movedColumnHeader) return;
      if (document.activeElement !== movedColumnHeader)
        movedColumnHeader.focus();
      updateFocusedColumnIndicator(movedColumnHeader);
    });
    finishColumnDrop();
  }
}

function shouldRenderExpandedRow(row: Row<TData>) {
  return (
    Boolean(props.getRowCanExpand) &&
    Boolean(row.getIsExpanded() || expandedRowsToRender.value[row.id])
  );
}

function handleExpandedRowAfterLeave(rowId: string) {
  const nextRowsToRender = { ...expandedRowsToRender.value };
  delete nextRowsToRender[rowId];
  expandedRowsToRender.value = nextRowsToRender;
}

function handleColumnDragOver(event: DragEvent, targetColumnId: string) {
  if (!props.draggableColumns || !draggingColumnId.value) return;

  event.preventDefault();
  const target = event.currentTarget as HTMLElement;
  const dropSide = getDropSide(target, event);
  activeDropTarget.value = { columnId: targetColumnId, side: dropSide };
  updateDropIndicator(target, dropSide);
  if (event.dataTransfer) event.dataTransfer.dropEffect = 'move';
}

function handleColumnDrop(event: DragEvent, targetColumnId: string) {
  if (!props.draggableColumns) return;

  event.preventDefault();

  const draggedColumnId = getDragData(event) ?? draggingColumnId.value;
  const target = event.currentTarget as HTMLElement;
  const dropSide =
    activeDropTarget.value?.columnId === targetColumnId
      ? activeDropTarget.value.side
      : getDropSide(target, event);

  if (draggedColumnId) {
    reorderColumn(draggedColumnId, targetColumnId, dropSide);
  }

  finishColumnDrop();
}

function handleColumnDragEnd() {
  finishColumnDrop();
}

function getColumnHeaderEventHandlers(header: Header<TData, unknown>) {
  if (!canDragHeader(header)) return {};

  return {
    dragstart: (event: DragEvent) => handleColumnDragStart(event, header),
    dragover: (event: DragEvent) =>
      handleColumnDragOver(event, header.column.id),
    drop: (event: DragEvent) => handleColumnDrop(event, header.column.id),
    dragend: handleColumnDragEnd,
    keydown: (event: KeyboardEvent) => handleColumnKeydown(event, header),
    focus: handleColumnHeaderFocus,
    click: handleColumnHeaderClick,
    blur: handleColumnHeaderBlur,
  };
}

function getColumnDropEventHandlers(columnId: string) {
  if (!props.draggableColumns) return {};

  return {
    dragover: (event: DragEvent) => handleColumnDragOver(event, columnId),
    drop: (event: DragEvent) => handleColumnDrop(event, columnId),
  };
}

function handleRowClick(event: MouseEvent | KeyboardEvent, original: TData) {
  if (!props.onRowClick) return;
  const target = event.target as HTMLElement;
  if (target.closest?.(INTERACTIVE_SELECTOR)) return;
  event.preventDefault();
  props.onRowClick(original);
}

defineExpose({
  table,
});
</script>

<template>
  <div
    ref="wrapperRef"
    class="relative"
    :style="
      props.tableMinHeight ? { minHeight: props.tableMinHeight } : undefined
    "
  >
    <div
      v-if="focusedColumnIndicator"
      data-column-focus-indicator
      class="pointer-events-none absolute z-[9] rounded-sm border-2 border-border-surface-focus-ring"
      :style="{
        left: `${focusedColumnIndicator.left}px`,
        top: `${focusedColumnIndicator.top}px`,
        width: `${focusedColumnIndicator.width}px`,
        height: `${focusedColumnIndicator.height}px`,
      }"
    />
    <div
      v-if="dropIndicator"
      data-column-drop-indicator
      class="pointer-events-none absolute z-10 bg-accent-default"
      :style="{
        left: `${dropIndicator.left}px`,
        top: `${dropIndicator.top}px`,
        width: `${DROP_LINE_WIDTH_PX}px`,
        height: `${dropIndicator.height}px`,
      }"
    />
    <UiTable>
      <UiTableHeader>
        <UiTableRow
          v-for="headerGroup in table.getHeaderGroups()"
          :key="headerGroup.id"
        >
          <UiTableHead
            v-for="header in headerGroup.headers"
            :key="header.id"
            :data-column-id="header.column.id"
            :colspan="header.colSpan"
            :draggable="canDragHeader(header)"
            :tabindex="canDragHeader(header) ? 0 : undefined"
            :class="[
              canDragHeader(header) ? 'select-none outline-none' : undefined,
              getColumnHeaderInteractionClass(header.column.id),
              getColumnHeaderDragStateClass(header.column.id),
            ]"
            :style="
              header.getSize() !== DEFAULT_COLUMN_SIZE
                ? { width: `${header.getSize()}px` }
                : undefined
            "
            v-on="getColumnHeaderEventHandlers(header)"
          >
            <FlexRender
              v-if="!header.isPlaceholder"
              :render="header.column.columnDef.header"
              :props="header.getContext()"
            />
          </UiTableHead>
        </UiTableRow>
      </UiTableHeader>
      <UiTableBody
        :class="props.tableMinHeight ? '[&_tr:last-child]:border-b' : undefined"
      >
        <template v-if="table.getRowModel().rows?.length">
          <template
            v-for="(row, rowIndex) in table.getRowModel().rows"
            :key="row.id"
          >
            <UiTableRow
              :data-state="row.getIsSelected() ? 'selected' : undefined"
              :class="props.onRowClick ? 'cursor-pointer' : undefined"
            >
              <UiTableCell
                v-for="cell in row.getVisibleCells()"
                :key="cell.id"
                :draggable="props.draggableColumns"
                :style="
                  cell.column.getSize() !== DEFAULT_COLUMN_SIZE
                    ? { width: `${cell.column.getSize()}px` }
                    : undefined
                "
                :class="
                  getColumnCellDragStateClass(
                    cell.column.id,
                    rowIndex === table.getRowModel().rows.length - 1,
                  )
                "
                v-on="getColumnDropEventHandlers(cell.column.id)"
                @dragstart="handleColumnCellDragStart($event, cell.column.id)"
                @dragend="handleColumnDragEnd"
                @click="handleRowClick($event, row.original)"
                @keydown.enter="handleRowClick($event, row.original)"
                @keydown.space="handleRowClick($event, row.original)"
              >
                <FlexRender
                  :render="cell.column.columnDef.cell"
                  :props="cell.getContext()"
                />
              </UiTableCell>
            </UiTableRow>
            <UiTableRow
              v-if="shouldRenderExpandedRow(row) && $slots['expanded-row']"
            >
              <UiTableCell :colspan="row.getVisibleCells().length" class="p-0">
                <Transition
                  :enter-active-class="heightTransitionClass"
                  :leave-active-class="heightTransitionClass"
                  @enter="onHeightTransitionEnter"
                  @after-enter="onHeightTransitionAfterEnter"
                  @before-leave="onHeightTransitionBeforeLeave"
                  @leave="onHeightTransitionLeave"
                  @after-leave="handleExpandedRowAfterLeave(row.id)"
                >
                  <div v-if="row.getIsExpanded()">
                    <div class="p-2">
                      <slot name="expanded-row" :row="row" :table="table" />
                    </div>
                  </div>
                </Transition>
              </UiTableCell>
            </UiTableRow>
          </template>
        </template>
        <template v-else>
          <UiTableRow>
            <UiTableCell
              :colspan="props.columns.length"
              class="h-24 text-center"
            >
              <slot name="empty">{{ props.emptyText }}</slot>
            </UiTableCell>
          </UiTableRow>
        </template>
      </UiTableBody>
    </UiTable>
  </div>
</template>
