import {
  Table,
  TableContainer,
  Tbody,
  Thead,
  Tr,
  Th,
  Td,
  chakra,
  Flex,
} from "@chakra-ui/react";
import type { UseQueryResult } from "@tanstack/react-query";
import {
  useReactTable,
  getCoreRowModel,
  getFilteredRowModel,
  ColumnDef,
  createColumnHelper,
  AccessorFn,
  DeepKeys,
  flexRender,
  PaginationState,
  OnChangeFn,
} from "@tanstack/react-table";
import { useMemo } from "react";

import { DataRowSkeleton } from "./DataRowSkeleton";
import { DataTableLoader } from "./DataTableLoader";
import { DataTablePagination } from "./DataTablePagination";
import {
  DataTableRecord,
  DataTableTransformer,
  PaginatedResponse,
  TData,
} from "../types";

export type ArrayElement<ArrayType extends readonly unknown[]> =
  ArrayType extends readonly (infer ElementType)[] ? ElementType : never;

export interface DataTableProps<T extends TData> {
  /**
   * Th title of the table
   */
  title: string;
  /**
   * The fields to display in the table
   */
  headerFields: (keyof T | "_")[];
  /**
   * The query response from the useDataTable hook
   */
  queryResponse: UseQueryResult<PaginatedResponse<T>>;
  /**
   * The pagination state from the useDataTable hook
   */
  pagination: PaginationState;
  /**
   * The pagination change handler from the useDataTable hook
   */
  onPaginationChange: OnChangeFn<PaginationState>;
  /**
   * An optional transform function to apply to the data used in the table
   */
  transformer?: DataTableTransformer<T>;
}

/**
 * A data table component that for formatted, paginated data
 *
 * @example
 * ```tsx
 * // Example usage
 * const dataTableProps = useDataTable<DataType>({
 *   queryFn: getPaginatedTestData,
 *   queryKey: "example-key",
 * })
 *
 * return <DataTable<DataType> {...dataTableProps} />
 * ```
 *
 * @param params The props for the data table
 * @returns A formatted data table
 */
export const DataTable = <T extends TData>({
  title,
  headerFields,
  queryResponse,
  pagination,
  onPaginationChange,
  transformer,
}: DataTableProps<T>) => {
  const data: T[] = queryResponse.data?.data ?? [];
  const pageCount = Math.floor(queryResponse.data?.pageCount ?? 0);

  const columnHelper = createColumnHelper<T>();

  const columns = useMemo<ColumnDef<DataTableRecord<T>>[]>(
    () => [
      {
        header: title,
        columns: (headerFields as string[]).map((field) =>
          columnHelper.accessor(field as AccessorFn<T, unknown> | DeepKeys<T>, {
            id: field,
            cell: (info) => info.getValue(),
            header: (info) =>
              info.column.id
                .replace(/([a-z])([A-Z])/g, "$1 $2")
                .toLocaleUpperCase(),
            footer: (info) => info.column.id,
          })
        ),
      },
    ],
    [title, columnHelper, headerFields]
  );

  const table = useReactTable<DataTableRecord<T>>({
    data: transformer ? transformer(data) : data,
    pageCount,
    columns,
    getCoreRowModel: getCoreRowModel<T>(),
    getFilteredRowModel: getFilteredRowModel<T>(),
    state: { pagination },
    onPaginationChange,
    manualPagination: true,
  });

  return (
    <DataTableContainer>
      <TableContainer w="full">
        <Table variant="simple" bg="white" rounded="lg">
          <Thead>
            {table.getHeaderGroups().map((headerGroup) => (
              <Tr key={headerGroup.id}>
                {headerGroup.headers.map((header) => (
                  <StyledTh key={header.id} colSpan={header.colSpan}>
                    {header.isPlaceholder ? null : (
                      <>
                        {flexRender(
                          header.column.columnDef.header,
                          header.getContext()
                        )}
                      </>
                    )}
                  </StyledTh>
                ))}
              </Tr>
            ))}
          </Thead>

          <Tbody>
            {table.getRowModel().rows.map((row) => (
              <StyledRow key={row.id}>
                {row.getVisibleCells().map((cell) => (
                  <StyledCell key={cell.id}>
                    {queryResponse.isFetching ? (
                      <DataRowSkeleton />
                    ) : (
                      flexRender(cell.column.columnDef.cell, cell.getContext())
                    )}
                  </StyledCell>
                ))}
              </StyledRow>
            ))}
          </Tbody>
        </Table>
      </TableContainer>

      <DataTableLoader
        isLoading={queryResponse.isFetching}
        refetch={queryResponse.refetch}
      />
      <PaginationContainer>
        <DataTablePagination table={table} />
      </PaginationContainer>
    </DataTableContainer>
  );
};

const DataTableContainer = chakra(Flex, {
  baseStyle: {
    flexDirection: "column",
    gap: "10px",
    boxShadow: "base",
    rounded: "lg",
    overflow: "hidden",
    position: "relative",
  },
});

const StyledTh = chakra(Th, {
  baseStyle: {
    fontFamily: "Roboto",
    color: "primaryAccent.500",
    bg: "gray.50",
  },
});

const StyledRow = chakra(Tr, {
  baseStyle: {
    fontSize: "sm",
  },
});

const StyledCell = chakra(Td, {
  baseStyle: {},
});

const PaginationContainer = chakra(Flex, {
  baseStyle: {
    gap: "5px",
    justifyContent: "flex-end",
    p: "10px",
    pt: "0px",
  },
});
