// ----- Imports -----

import {
  Box,
  Checkbox,
  HStack,
  Spacer as ChakraSpacer,
  Table as ChakraTable,
  Tbody,
  Td,
  Text,
  Th,
  Thead,
  Tr,
  VStack,
} from "@chakra-ui/react";
import { Pagination as UIKitPagination } from "../pagination/Pagination.jsx";
import { toFaDigits } from "../utils/numbers.js";
import {
  flexRender,
  getCoreRowModel,
  useReactTable,
} from "@tanstack/react-table";
import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from "react";

// ----- Context -----

const DataTableContext = createContext(null);

function useDataTable() {
  const context = useContext(DataTableContext);

  if (!context) {
    throw new Error("DataTable components must be used inside <DataTable>");
  }

  return context;
}

// ----- Root Component -----

function Root({
  children,

  data = [],
  columns = [],

  page = 1,
  size = 10,

  pagination,
  onPageChange,

  getRowId = (row) => String(row?.id),

  enableSelection = false,
  showIndex = true,

  defaultSelectedIds = [],
  onSelectionChange,

  ...props
}) {
  // ----- Normalized Pagination Props -----

  const normalizedPage = Number(page) || 1;
  const normalizedSize = Number(size) || 10;

  const totalCount = pagination?.total ?? pagination?.totalCount ?? 0;

  const totalPages =
    pagination?.total_pages ??
    pagination?.totalPages ??
    (totalCount > 0 ? Math.ceil(totalCount / normalizedSize) : 0);

  // ----- TanStack Pagination State -----

  const paginationState = useMemo(() => {
    return {
      pageIndex: Math.max(normalizedPage - 1, 0),
      pageSize: normalizedSize,
    };
  }, [normalizedPage, normalizedSize]);

  const handlePaginationChange = useCallback(
    (updater) => {
      const next =
        typeof updater === "function" ? updater(paginationState) : updater;

      const nextPage = next.pageIndex + 1;

      if (nextPage !== normalizedPage) {
        onPageChange?.(nextPage);
      }
    },
    [paginationState, normalizedPage, onPageChange],
  );

  // ----- Selection State -----

  const [selectedIds, setSelectedIdsState] = useState(defaultSelectedIds);

  const setSelectedIds = useCallback(
    (updater) => {
      setSelectedIdsState((prev) => {
        const next = typeof updater === "function" ? updater(prev) : updater;

        onSelectionChange?.(next);

        return next;
      });
    },
    [onSelectionChange],
  );

  // ----- Reset Selection On Page Change -----

  const previousPageRef = useRef({
    page: normalizedPage,
    size: normalizedSize,
  });

  useEffect(() => {
    if (!enableSelection) {
      return;
    }

    const previous = previousPageRef.current;

    const pageChanged =
      previous.page !== normalizedPage || previous.size !== normalizedSize;

    previousPageRef.current = {
      page: normalizedPage,
      size: normalizedSize,
    };

    if (pageChanged) {
      setSelectedIds([]);
    }
  }, [enableSelection, normalizedPage, normalizedSize, setSelectedIds]);

  // ----- Page Row Ids -----

  const allPageIds = useMemo(() => {
    return data
      .map((item) => getRowId(item))
      .filter((id) => id !== undefined && id !== null && id !== "");
  }, [data, getRowId]);

  const pageSelectedIds = useMemo(() => {
    return selectedIds.filter((id) => allPageIds.includes(id));
  }, [allPageIds, selectedIds]);

  // ----- Selection Status -----

  const isAllSelected =
    allPageIds.length > 0 && pageSelectedIds.length === allPageIds.length;

  const isIndeterminate =
    pageSelectedIds.length > 0 && pageSelectedIds.length < allPageIds.length;

  // ----- Selection Handlers -----

  const clearSelection = useCallback(() => {
    setSelectedIds([]);
  }, [setSelectedIds]);

  const onSelectionChangedId = useCallback(
    (id) => {
      setSelectedIds((prev) => {
        if (!allPageIds.includes(id)) {
          return prev;
        }

        if (prev.includes(id)) {
          return prev.filter((item) => item !== id);
        }

        return [...prev, id];
      });
    },
    [allPageIds, setSelectedIds],
  );

  const toggleRowSelection = onSelectionChangedId;

  const selectAll = useCallback(() => {
    setSelectedIds(allPageIds);
  }, [allPageIds, setSelectedIds]);

  const unselectAll = useCallback(() => {
    setSelectedIds([]);
  }, [setSelectedIds]);

  const toggleSelectAll = useCallback(
    (checked) => {
      if (checked) {
        selectAll();
        return;
      }

      unselectAll();
    },
    [selectAll, unselectAll],
  );

  // ----- Table Columns -----

  const tableColumns = useMemo(() => {
    const result = [];

    if (enableSelection) {
      result.push({
        id: "__selection",
        header: ({ table }) => {
          const meta = table.options.meta;

          return (
            <Checkbox
              isChecked={meta.isAllSelected}
              isIndeterminate={meta.isIndeterminate}
              onChange={(e) => meta.toggleSelectAll(e.target.checked)}
              colorScheme="purple"
              size="md"
            />
          );
        },
        cell: ({ row, table }) => {
          const meta = table.options.meta;
          const rowId = meta.getRowId(row.original);

          return (
            <Checkbox
              isChecked={meta.selectedIds.includes(rowId)}
              onChange={() => meta.onSelectionChangedId(rowId)}
              colorScheme="purple"
              size="md"
            />
          );
        },
        meta: { isCentered: true },
      });
    }

    if (showIndex) {
      result.push({
        id: "__index",
        header: "ردیف",
        cell: ({ row, table }) => {
          const { pageIndex, pageSize } = table.getState().pagination;
          const rowNumber = pageIndex * pageSize + row.index + 1;

          return (
            <Text fontWeight="bold" sx={{ fontVariantNumeric: "normal" }}>
              {toFaDigits(rowNumber)}
            </Text>
          );
        },
        meta: { isCentered: true },
      });
    }

    return [...result, ...columns];
  }, [columns, enableSelection, showIndex]);

  // ----- Table Instance -----

  const table = useReactTable({
    data,
    columns: tableColumns,
    getCoreRowModel: getCoreRowModel(),
    getRowId,

    manualPagination: true,
    pageCount: totalPages,
    rowCount: totalCount,

    state: {
      pagination: paginationState,
    },

    onPaginationChange: handlePaginationChange,

    meta: {
      selectedIds,
      onSelectionChangedId,

      getRowId,

      allPageIds,
      pageSelectedIds,

      isAllSelected,
      isIndeterminate,

      selectAll,
      unselectAll,
      toggleSelectAll,
      clearSelection,

      totalCount,
      totalPages,
    },
  });

  // ----- Provider Value -----

  const value = {
    table,

    data,
    columns: tableColumns,

    page: normalizedPage,
    size: normalizedSize,

    pagination,
    onPageChange,

    paginationState,
    totalCount,
    totalPages,

    getRowId,

    enableSelection,

    selectedIds,
    setSelectedIds,
    clearSelection,

    allPageIds,
    pageSelectedIds,

    isAllSelected,
    isIndeterminate,

    selectAll,
    unselectAll,
    toggleSelectAll,

    onSelectionChangedId,
    toggleRowSelection,
  };

  // ----- Root Render -----

  return (
    <DataTableContext.Provider value={value}>
      <VStack spacing={4} w="full" {...props}>
        {children}
      </VStack>
    </DataTableContext.Provider>
  );
}

// ----- Topbar Components -----

function Topbar({ children, ...props }) {
  return (
    <HStack w="full" px={2} spacing={4} {...props}>
      {children}
    </HStack>
  );
}

function Title({ children, ...props }) {
  return (
    <Text fontSize="lg" textAlign="start" fontWeight="bold" {...props}>
      {children}
    </Text>
  );
}

function Spacer(props) {
  return <ChakraSpacer {...props} />;
}

function Actions({ children, ...props }) {
  const context = useDataTable();

  return (
    <HStack spacing={4} {...props}>
      {typeof children === "function" ? children(context) : children}
    </HStack>
  );
}

// ----- Selection Components -----

function SelectionActions({ children }) {
  const context = useDataTable();

  if (context.selectedIds.length === 0) {
    return null;
  }

  return typeof children === "function" ? children(context) : children;
}

function SelectAll({ children = "انتخاب همه", ...props }) {
  const {
    data,
    enableSelection,
    isAllSelected,
    isIndeterminate,
    toggleSelectAll,
  } = useDataTable();

  if (!enableSelection || data.length === 0) {
    return null;
  }

  return (
    <Checkbox
      isChecked={isAllSelected}
      isIndeterminate={isIndeterminate}
      onChange={(e) => toggleSelectAll(e.target.checked)}
      colorScheme="purple"
      size="lg"
      {...props}
    >
      {children}
    </Checkbox>
  );
}

function SelectedCount({ children, ...props }) {
  const { selectedIds } = useDataTable();

  if (selectedIds.length === 0) {
    return null;
  }

  if (typeof children === "function") {
    return children({ count: selectedIds.length, selectedIds });
  }

  return (
    <Text fontSize="sm" color="gray.600" {...props}>
      {toFaDigits(selectedIds.length)} مورد انتخاب شده
    </Text>
  );
}

// ----- Content Wrapper -----

function Content({ children, ...props }) {
  return (
    <Box
      display="flex"
      flexDirection="column"
      w="full"
      h="100%"
      minH={0}
      gap={4}
      {...props}
    >
      {children}
    </Box>
  );
}

// ----- Table Container -----

function Table({
  children,
  minW = "1200px",
  heightOffset = 355,
  containerProps,
  tableProps,
}) {
  return (
    <Box
      flex="1"
      minH={`calc(100vh - ${heightOffset}px)`}
      maxH={`calc(100vh - ${heightOffset}px)`}
      overflowX="auto"
      overflowY="auto"
      bg="white"
      borderRadius="lg"
      borderWidth="1px"
      borderColor="gray.200"
      sx={{
        "::-webkit-scrollbar": {
          width: "6px",
          height: "6px",
        },
        "::-webkit-scrollbar-thumb": {
          background: "#CBD5E0",
          borderRadius: "8px",
        },
        "::-webkit-scrollbar-track": {
          background: "transparent",
        },
      }}
      {...containerProps}
    >
      <ChakraTable
        variant="simple"
        minW={minW}
        sx={{
          "th, td": {
            borderColor: "gray.200",
            whiteSpace: "nowrap",
            px: 2,
          },
          "tbody tr:last-child td": {
            borderBottom: "none",
          },
        }}
        {...tableProps}
      >
        {children ?? (
          <>
            <DataTable.Head />
            <DataTable.Body />
          </>
        )}
      </ChakraTable>
    </Box>
  );
}

// ----- Table Head -----

function Head() {
  const { table } = useDataTable();

  return (
    <Thead position="sticky" top={0} bg="white" zIndex={1}>
      {table.getHeaderGroups().map((group) => (
        <Tr key={group.id} borderBottomWidth="1px" borderColor="gray.200">
          {group.headers.map((header) => (
            <Th
              key={header.id}
              py={4}
              borderRightWidth="1px"
              borderRightColor="gray.200"
              fontSize="md"
              color="gray.600"
              textAlign={
                header.column.columnDef.meta?.isCentered ? "center" : "start"
              }
              verticalAlign="middle"
              textTransform="none"
            >
              {header.isPlaceholder
                ? null
                : flexRender(
                    header.column.columnDef.header,
                    header.getContext(),
                  )}
            </Th>
          ))}
        </Tr>
      ))}
    </Thead>
  );
}

// ----- Table Body -----

function Body({ emptyText = "موردی وجود ندارد" }) {
  const { table, columns } = useDataTable();

  const rows = table.getRowModel().rows;

  return (
    <Tbody>
      {rows.length > 0 ? (
        rows.map((row) => (
          <Tr
            key={row.id}
            borderBottomWidth="1px"
            borderColor="gray.200"
            _hover={{ bg: "blue.50" }}
            transition="0.2s"
          >
            {row.getVisibleCells().map((cell) => (
              <Td
                key={cell.id}
                py={3}
                fontSize="md"
                color="gray.700"
                borderRightWidth="1px"
                borderRightColor="gray.200"
                textAlign={
                  cell.column.columnDef.meta?.isCentered ? "center" : "start"
                }
                verticalAlign="middle"
              >
                {cell.column.columnDef.cell
                  ? flexRender(cell.column.columnDef.cell, cell.getContext())
                  : cell.getValue()}
              </Td>
            ))}
          </Tr>
        ))
      ) : (
        <Tr>
          <Td colSpan={columns.length} py={8} textAlign="center">
            <Text color="gray.500">{emptyText}</Text>
          </Td>
        </Tr>
      )}
    </Tbody>
  );
}

// ----- Pagination -----

function Pagination({
  pagination: paginationProp,
  onPageChange: onPageChangeProp,
}) {
  const context = useDataTable();

  const { table } = context;

  const pagination = paginationProp ?? context.pagination;

  const totalCount =
    pagination?.total ?? pagination?.totalCount ?? context.totalCount ?? 0;

  const totalPages =
    pagination?.total_pages ??
    pagination?.totalPages ??
    context.totalPages ??
    table.getPageCount() ??
    0;

  const currentPage = table.getState().pagination.pageIndex + 1;

  const handlePageChange = (nextPage) => {
    if (onPageChangeProp) {
      onPageChangeProp(nextPage);
      return;
    }

    table.setPageIndex(nextPage - 1);
  };

  if (!context.onPageChange && !onPageChangeProp) {
    return null;
  }

  if (totalPages <= 0) {
    return null;
  }

  return (
    <Box mt={1} flexShrink={0}>
      <UIKitPagination
        currentPage={currentPage}
        onPageChange={handlePageChange}
        totalCount={totalCount}
        totalPages={totalPages}
      />
    </Box>
  );
}

// ----- Compound Export -----

export const DataTable = Object.assign(Root, {
  Topbar,
  Title,
  Spacer,
  Actions,
  SelectionActions,
  SelectAll,
  SelectedCount,
  Content,
  Table,
  Head,
  Body,
  Pagination,
});