import {
  AccessorFnColumnDef,
  ColumnDef,
  DisplayColumnDef,
} from "@tanstack/react-table";
import * as React from "react";
import { SortIcon } from "../../icons";
import {
  useCollection,
  useShallowRender,
  type BaseCollection,
  type Node,
} from "../Collection";
import { Checkbox, DataTable as DataTableInner } from "./DataTableImpl";

interface DataTableProps<TData> extends React.ComponentProps<"div"> {
  "data-id": string;
  data: TData[];
  children?: React.ReactNode;
  __length?: number;
  searchTerm?: string;
  initialPageSize?: number;
}

const EMPTY_DATA = [];
export function DataTable<TData, TValue>({
  data = [],
  children,
  __length,
  ...props
}: DataTableProps<TData>) {
  const { portal, collection } = useCollection({ children });
  return (
    <>
      {portal}
      <DataTableComponent
        data={data || EMPTY_DATA}
        collection={collection}
        {...props}
      />
    </>
  );
}

type DataTableComponentProps<TData> = DataTableProps<TData> & {
  collection: BaseCollection<object>;
  "data-id": string;
  data: TData[];
  searchTerm?: string;
};

function DataTableComponent<TData, TValue>({
  data,
  collection,
  ...props
}: DataTableComponentProps<TData>) {
  const selectAbleRef = React.useRef<boolean>(false);
  const cache = React.useMemo(
    () => new WeakMap<Node<object>, ColumnDef<TData, TValue>>(),
    [],
  );
  const columns = React.useMemo(() => {
    selectAbleRef.current = false;
    const columnDefs: Array<ColumnDef<TData, TValue> & { className?: string }> =
      [];
    for (const node of collection) {
      let columnDef = cache.get(node);
      if (!columnDef) {
        switch (node.type) {
          case "value":
            columnDef = {
              id: node.props.header,
              accessorFn: node.props.accessorFn,
              meta: {
                className: node.props.className,
              },
              header:
                node.props.enableSorting === false
                  ? node.props.header
                  : ({ column }) => (
                      <button
                        data-sortable
                        data-active={column.getIsSorted() !== false}
                        onClick={() =>
                          column.toggleSorting(column.getIsSorted() === "asc")
                        }
                      >
                        {node.props.header}
                        <SortIcon direction={column.getIsSorted() || "none"} />
                      </button>
                    ),
              cell: ({ row: { original } }) => node.props.children?.(original),
            };
            break;
          case "display":
            columnDef = {
              id: node.props.header,
              meta: {
                className: node.props.className,
              },
              accessorFn: node.props.accessorFn,
              header: node.props.header,
              cell: ({ row: { original } }) => node.props.children?.(original),
            };
            break;
          case "checkbox":
            selectAbleRef.current = true;
            columnDef = {
              id: node.props.id,
              header: ({ table }) => (
                <Checkbox
                  checked={
                    table.getIsAllPageRowsSelected() ||
                    (table.getIsSomePageRowsSelected() && "indeterminate")
                  }
                  onCheckedChange={(value) =>
                    table.toggleAllPageRowsSelected(!!value)
                  }
                  aria-label="Select all"
                />
              ),
              cell: ({ row }) => (
                <Checkbox
                  checked={row.getIsSelected()}
                  onCheckedChange={(value) => row.toggleSelected(!!value)}
                  aria-label="Select row"
                />
              ),
              enableSorting: false,
              enableHiding: false,
            };
            break;
          default:
            console.log("unknown child node", node);
            break;
        }
        if (columnDef) {
          cache.set(node, columnDef);
        }
      }
      if (columnDef) {
        columnDefs.push(columnDef);
      }
    }
    return columnDefs;
  }, [collection, cache]);
  return (
    <DataTableInner
      data={data || EMPTY_DATA}
      columns={columns}
      selectable={selectAbleRef.current}
      {...props}
    />
  );
}

type ValueProps<TData, TValue> = Omit<
  AccessorFnColumnDef<TData, TValue>,
  "cell"
> & {
  children: (origional: TData) => React.ReactNode;
  className?: string;
};

function ValueColumn<TData, TValue = unknown>(
  props: ValueProps<TData, TValue> & {
    ref?: React.ForwardedRef<Element>;
  },
) {
  return useShallowRender("value", props, props.ref!);
}

DataTable.ValueColumn = ValueColumn;

type DisplayProps<TData, TValue> = Omit<
  DisplayColumnDef<TData, TValue>,
  "cell"
> & {
  children: (origional: TData) => React.ReactNode;
  className?: string;
};

function DisplayColumn<TData, TValue = unknown>(
  props: DisplayProps<TData, TValue> & {
    ref?: React.ForwardedRef<Element>;
  },
) {
  return useShallowRender("display", props, props.ref!);
}

DataTable.DisplayColumn = DisplayColumn;

function CheckboxColumn(props: {
  id: string;
  className?: string;
  ref?: React.ForwardedRef<Element>;
}) {
  return useShallowRender("checkbox", props, props.ref!);
}

DataTable.CheckboxColumn = CheckboxColumn;