import { useEffect, useState } from "react";
import { useQuery, useQueryClient } from "@tanstack/react-query";
import isNumber from "lodash/isNumber";

export type PaginationDirection = "ASC" | "DESC";
export interface PaginationQueryParameters {
  cursor: string | null;
  cutoff: number;
  limit: number;
  offset: number;
  direction: PaginationDirection;
}
export interface CountQueryParameters {
  cutoff: number;
}
export type CursorPaginatedItem<T> = T & { cursor: string };

interface PaginationReturns<T> {
  // current page viewed by the user
  currentPage: number;
  // same as currentPage but starting from 0
  currentPageIndex: number;
  // the total number of pages
  pageCount: number;
  // the total number of rows
  rowCount: number;
  // is data loading
  isLoading: boolean;
  // the current page rows
  data: T[];
  // use this function to change the page the user is viewing
  changePage: (page: number) => any;
  // same as changePage but takes an index
  changePageIndex: (pageIndex: number) => any;
  // scrap all the internal states and start fetching the first page
  reset: () => any;
}

export function usePagination<T extends { cursor: string | null }>(
  queryKey: unknown | unknown[],
  // should return one page given the parameters
  query: (params: PaginationQueryParameters) => Promise<T[]>,
  // Should return the total number of rows in the dataset
  countQuery: (parameters: CountQueryParameters) => Promise<number>,
  limit: number,
  direction: "ASC" | "DESC",
  // Put any dependencies you are using in query or countQuery, any changes will reset the pagination
  deps?: React.DependencyList | undefined
): PaginationReturns<T> {
  const [currentPageIndex, setCurrentPageIndex] = useState(0);
  const [cutoff, setCutoff] = useState(0);

  const queryClient = useQueryClient();

  const { data: rowCount, isLoading: isCountLoading } = useQuery({
    enabled: cutoff !== 0,
    queryKey: [queryKey, "count", { cutoff }, deps],
    keepPreviousData: true,
    queryFn: () => {
      return countQuery({
        cutoff,
      });
    },
  });

  const pageCount =
    rowCount === 0 || rowCount === undefined ? 1 : Math.ceil(rowCount / limit);

  const { data, isLoading } = useQuery({
    // Don't run until the count was loaded or if we are requesting the first page
    // It's safe to load the first page as we are not passing a cursor to get it.
    enabled: currentPageIndex === 1 || !isCountLoading,
    queryKey: [
      queryKey,
      "pages",
      { limit, direction, cutoff },
      deps,
      currentPageIndex,
    ],
    queryFn: async () => {
      const pageMap = extractPageDataFromCache([
        queryKey,
        "pages",
        { limit, direction, cutoff },
        deps,
        currentPageIndex,
      ]);

      const params = {
        ...getParams(currentPageIndex, pageCount, pageMap, direction, limit),
        cutoff,
      };

      const data = await query(params);

      return transformData(data, params);
    },
    keepPreviousData: true,
  });

  function transformData(d: T[], params: PaginationQueryParameters): T[] {
    let data = Array.from(d);

    const isQueryingFromEnd = params.direction !== direction;
    const isLastPage = isQueryingFromEnd && params.cursor === null;
    const isLastPageFull = (rowCount ?? 0) % limit === 0;

    if (isLastPage && !isLastPageFull) {
      // get only the few elements that bellong on the last page
      data = data.slice(0, (rowCount ?? 0) % limit);
    }

    if (isQueryingFromEnd) {
      // the data start from the end so we want to revert it
      data.reverse();
    }

    return data;
  }

  function extractPageDataFromCache(queryKey: unknown[]): Map<number, T[]> {
    if (!isNumber(queryKey[queryKey.length - 1])) {
      throw new Error("page number should always be last in queryKey");
    }

    const cachedPageData = queryClient.getQueryCache().findAll({
      queryKey: queryKey.slice(0, -1),
    });

    const pageTuples = cachedPageData
      .map((x) => {
        const pageNumber = x.queryKey[queryKey.length - 1];
        const data = x.state.data;

        if (!isNumber(pageNumber) || !Array.isArray(data)) {
          return null;
        }

        return [pageNumber, data] as const;
      })
      .filter((x) => Boolean(x)) as [number, any[]][];

    return new Map<number, T[]>(pageTuples);
  }

  function changePageIndex(pageIndex: number) {
    setCurrentPageIndex(pageIndex);
  }

  useEffect(() => {
    reset();
  }, [...(deps ?? []), limit, direction]);

  function reset() {
    setCutoff(Date.now());
    setCurrentPageIndex(0);
  }

  return {
    currentPage: currentPageIndex + 1,
    currentPageIndex: currentPageIndex,
    pageCount,
    rowCount: rowCount ?? 0,
    isLoading,
    data: data ?? [],
    changePageIndex,
    changePage(page: number) {
      changePageIndex(page - 1);
    },
    reset,
  };
}
export function useOffsetPagination<T extends { cursor: string | null }>(
  queryKey: unknown | unknown[],
  // should return one page given the parameters
  query: (params: PaginationQueryParameters) => Promise<T[]>,
  // Should return the total number of rows in the dataset
  countQuery: (parameters: CountQueryParameters) => Promise<number>,
  limit: number,
  direction: "ASC" | "DESC",
  // Put any dependencies you are using in query or countQuery, any changes will reset the pagination
  deps?: React.DependencyList | undefined
): PaginationReturns<T> {
  const [currentPageIndex, setCurrentPageIndex] = useState(0);
  const [cutoff, setCutoff] = useState(0);

  const queryClient = useQueryClient();

  const { data: rowCount, isLoading: isCountLoading } = useQuery({
    enabled: cutoff !== 0,
    queryKey: [queryKey, "count", { cutoff }, deps],
    keepPreviousData: true,
    queryFn: () => {
      return countQuery({
        cutoff,
      });
    },
  });

  const pageCount =
    rowCount === 0 || rowCount === undefined ? 1 : Math.ceil(rowCount / limit);

  const { data, isLoading } = useQuery({
    // Don't run until the count was loaded or if we are requesting the first page
    // It's safe to load the first page as we are not passing a cursor to get it.
    enabled: currentPageIndex === 0 || !isCountLoading,
    queryKey: [
      queryKey,
      "pages",
      { limit, direction, cutoff },
      deps,
      currentPageIndex,
    ],
    queryFn: async () => {
      const pageMap = extractPageDataFromCache([
        queryKey,
        "pages",
        { limit, direction, cutoff },
        deps,
        currentPageIndex,
      ]);

      const params = {
        ...getOffsetParams(
          currentPageIndex,
          pageCount,
          pageMap,
          direction,
          limit
        ),
        cutoff,
      };

      const data = await query(params);

      return transformData(data, params);
    },
    keepPreviousData: true,
  });

  function transformData(d: T[], params: PaginationQueryParameters): T[] {
    let data = Array.from(d);

    const isQueryingFromEnd = params.direction !== direction;
    const isLastPage = isQueryingFromEnd && params.cursor === null;
    const isLastPageFull = (rowCount ?? 0) % limit === 0;

    if (isLastPage && !isLastPageFull) {
      // get only the few elements that bellong on the last page
      data = data.slice(0, (rowCount ?? 0) % limit);
    }

    if (isQueryingFromEnd) {
      // the data start from the end so we want to revert it
      data.reverse();
    }

    return data;
  }

  function extractPageDataFromCache(queryKey: unknown[]): Map<number, T[]> {
    if (!isNumber(queryKey[queryKey.length - 1])) {
      throw new Error("page number should always be last in queryKey");
    }

    const cachedPageData = queryClient.getQueryCache().findAll({
      queryKey: queryKey.slice(0, -1),
    });

    const pageTuples = cachedPageData
      .map((x) => {
        const pageNumber = x.queryKey[queryKey.length - 1];
        const data = x.state.data;

        if (!isNumber(pageNumber) || !Array.isArray(data)) {
          return null;
        }

        return [pageNumber, data] as const;
      })
      .filter((x) => Boolean(x)) as [number, any[]][];

    return new Map<number, T[]>(pageTuples);
  }

  function changePageIndex(pageIndex: number) {
    setCurrentPageIndex(pageIndex);
  }

  useEffect(() => {
    reset();
  }, [...(deps ?? []), limit, direction]);

  function reset() {
    setCutoff(Date.now());
    setCurrentPageIndex(0);
  }

  return {
    currentPage: currentPageIndex + 1,
    currentPageIndex: currentPageIndex,
    pageCount,
    rowCount: rowCount ?? 0,
    isLoading,
    data: data ?? [],
    changePageIndex,
    changePage(page: number) {
      changePageIndex(page - 1);
    },
    reset,
  };
}

// Cursor based strategy
function getOffsetParams<T extends { cursor: string }>(
  pageIndex: number,
  pageCount: number,
  pageMap: Map<number, T[]>,
  direction: "ASC" | "DESC",
  limit: number
): {
  cursor: string | null;
  limit: number;
  offset: number;
  direction: "ASC" | "DESC";
} {
  return {
    cursor: null,
    limit,
    offset: limit * pageIndex,
    direction,
  };
}

// Cursor based strategy
function getParams<T extends { cursor: string }>(
  pageIndex: number,
  pageCount: number,
  pageMap: Map<number, T[]>,
  direction: "ASC" | "DESC",
  limit: number
): {
  cursor: string | null;
  limit: number;
  offset: number;
  direction: "ASC" | "DESC";
} {
  const closestPageIndex = getClosestPageIndex(pageIndex, pageCount, pageMap);
  const shouldInverseDirection = getShouldInverseDirection(
    closestPageIndex,
    pageIndex
  );
  const cursor = getCursor<T>(
    closestPageIndex,
    pageMap,
    shouldInverseDirection
  );
  const offset = getOffset(
    closestPageIndex,
    shouldInverseDirection,
    pageCount,
    pageIndex,
    limit
  );

  return {
    cursor,
    limit,
    offset,
    direction: shouldInverseDirection ? inverseDirection(direction) : direction,
  };
}

function inverseDirection(direction: "ASC" | "DESC"): "ASC" | "DESC" {
  if (direction === "ASC") {
    return "DESC";
  }

  return "ASC";
}

/*
 * We try to find the closest page to make sure we use the lowest offset
 * Most of the time the closest page is already loaded, if it's not that function
 * will return null which will result in the cursor being null
 */
function getClosestPageIndex<T extends { cursor: string }>(
  pageIndex: number,
  pageCount: number,
  pageMap: Map<number, T[]>
): null | number {
  let diff = 1;

  while (diff <= pageCount) {
    if (pageIndex - diff < 0) {
      return null;
    }

    if (pageIndex + diff > pageCount - 1) {
      return null;
    }

    if (pageMap.has(pageIndex - diff)) {
      return pageIndex - diff;
    }

    if (pageMap.has(pageIndex + diff)) {
      return pageIndex + diff;
    }

    diff++;
  }

  return null;
}

/*
 * We want to inverse the direction if we query from the end
 */
function getShouldInverseDirection(
  closestPageIndex: number | null,
  pageIndex: number
): boolean {
  if (closestPageIndex !== null) {
    return closestPageIndex > pageIndex;
  }

  if (pageIndex === 0) {
    return false;
  }

  return true;
}

function getCursor<T extends { cursor: string }>(
  closestPageIndex: number | null,
  pageMap: Map<number, T[]>,
  shouldInverseDirection: boolean
) {
  // closestPageIndex is null if we couldn't find a page to anchor the query from
  if (closestPageIndex === null) {
    return null;
  }

  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  const data = pageMap.get(closestPageIndex)!;

  return shouldInverseDirection ? data[0].cursor : data[data.length - 1].cursor;
}

function getOffset(
  closestPageIndex: number | null,
  shouldInverseDirection: boolean,
  pageCount: number,
  pageIndex: number,
  limit: number
) {
  if (closestPageIndex === null) {
    const edgePageIndex = shouldInverseDirection ? pageCount - 1 : 0;

    return Math.abs(edgePageIndex - pageIndex) * limit;
  }

  return (Math.abs(closestPageIndex - pageIndex) - 1) * limit;
}
