import { atom, useAtom, useAtomValue } from 'jotai';
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { MatrixClient, MatrixError, Room } from 'matrix-js-sdk';
import { IHierarchyRoom } from 'matrix-js-sdk/lib/@types/spaces';
import { QueryFunction, useInfiniteQuery } from '@tanstack/react-query';
import { useMatrixClient } from './useMatrixClient';
import { roomToParentsAtom } from '../state/room/roomToParents';
import { MSpaceChildContent, RoomType, StateEvent } from '../../types/matrix/room';
import { getAllParents, getStateEvents, isSpace, isValidChild } from '../utils/room';
import { isRoomId } from '../utils/matrix';
import { SortFunc, byOrderKey, byTsOldToNew, factoryRoomIdByActivity } from '../utils/sort';
import { useStateEventCallback } from './useStateEventCallback';
import { ErrorCode } from '../cs-errorcode';

export type HierarchyItemSpace = {
  roomId: string;
  content: MSpaceChildContent;
  ts: number;
  space: true;
  parentId?: string;
};

export type HierarchyItemRoom = {
  roomId: string;
  content: MSpaceChildContent;
  ts: number;
  parentId: string;
};

export type HierarchyItem = HierarchyItemSpace | HierarchyItemRoom;

type GetRoomCallback = (roomId: string) => Room | undefined;

const hierarchyItemTs: SortFunc<HierarchyItem> = (a, b) => byTsOldToNew(a.ts, b.ts);
const hierarchyItemByOrder: SortFunc<HierarchyItem> = (a, b) =>
  byOrderKey(a.content.order, b.content.order);

// we will paginate until 5000 items
const PER_PAGE_COUNT = 100;
const MAX_AUTO_PAGE_COUNT = 50;

export function buildParentByRoomId(rooms: Iterable<IHierarchyRoom>): Map<string, string> {
  const parentByRoomId = new Map<string, string>();
  for (const room of rooms) {
    if (!room?.room_id || !Array.isArray(room.children_state)) continue;
    for (const childState of room.children_state) {
      if (childState?.type !== StateEvent.SpaceChild) continue;
      if (typeof childState.state_key !== 'string') continue;
      parentByRoomId.set(childState.state_key, room.room_id);
    }
  }
  return parentByRoomId;
}

/** Space or forum container in hierarchy API (room_type may be absent on some servers). */
export function isHierarchySpaceRoom(room: IHierarchyRoom | undefined): boolean {
  if (!room) return false;
  if (room.room_type === RoomType.Space || room.room_type === RoomType.Forum) return true;
  if (Array.isArray(room.children_state)) {
    return room.children_state.some((c) => c?.type === StateEvent.SpaceChild);
  }
  return false;
}

export async function fetchAllHierarchyRooms(
  mx: MatrixClient,
  spaceId: string,
  maxDepth = 3
): Promise<IHierarchyRoom[]> {
  const allRooms: IHierarchyRoom[] = [];
  let nextBatch: string | undefined;

  for (let page = 0; page < MAX_AUTO_PAGE_COUNT; page += 1) {
    const result = await mx.getRoomHierarchy(spaceId, PER_PAGE_COUNT, maxDepth, false, nextBatch);
    allRooms.push(...result.rooms);
    nextBatch = result.next_batch;
    if (!nextBatch) break;
  }

  return allRooms;
}

const isChildSpace = (
  childId: string,
  getRoom: GetRoomCallback,
  spaceRooms: Set<string>,
  hierarchySummaries: Map<string, IHierarchyRoom>
): boolean =>
  isSpace(getRoom(childId) ?? null) ||
  spaceRooms.has(childId) ||
  isHierarchySpaceRoom(hierarchySummaries.get(childId));

const getChildRoomsFromSummaries = (
  spaceItemId: string,
  hierarchySummaries: Map<string, IHierarchyRoom>,
  parentByRoomId: Map<string, string>,
  spaceRooms: Set<string>,
  getRoom: GetRoomCallback
): HierarchyItemRoom[] => {
  const childItems: HierarchyItemRoom[] = [];
  hierarchySummaries.forEach((summary, roomId) => {
    if (parentByRoomId.get(roomId) !== spaceItemId) return;
    if (isChildSpace(roomId, getRoom, spaceRooms, hierarchySummaries)) return;

    childItems.push({
      roomId,
      content: { via: summary.via ?? [] },
      ts: 0,
      parentId: spaceItemId,
    });
  });
  return childItems;
};

const getHierarchySpaces = (
  rootSpaceId: string,
  getRoom: GetRoomCallback,
  spaceRooms: Set<string>,
  hierarchySummaries: Map<string, IHierarchyRoom>
): HierarchyItemSpace[] => {
  const rootSpaceItem: HierarchyItemSpace = {
    roomId: rootSpaceId,
    content: { via: [] },
    ts: 0,
    space: true,
  };
  let spaceItems: HierarchyItemSpace[] = [];
  const parentByRoomId = buildParentByRoomId(hierarchySummaries.values());

  const findAndCollectHierarchySpaces = (spaceItem: HierarchyItemSpace) => {
    if (spaceItems.find((item) => item.roomId === spaceItem.roomId)) return;
    const space = getRoom(spaceItem.roomId);
    spaceItems.push(spaceItem);

    if (space) {
      const childEvents = getStateEvents(space, StateEvent.SpaceChild);

      childEvents.forEach((childEvent) => {
        if (!isValidChild(childEvent)) return;
        const childId = childEvent.getStateKey();
        if (!childId || !isRoomId(childId)) return;

        if (isChildSpace(childId, getRoom, spaceRooms, hierarchySummaries)) {
          const childItem: HierarchyItemSpace = {
            roomId: childId,
            content: childEvent.getContent<MSpaceChildContent>(),
            ts: childEvent.getTs(),
            space: true,
            parentId: spaceItem.roomId,
          };
          findAndCollectHierarchySpaces(childItem);
        }
      });
    }

    // Unjoined sub-spaces: discover children from hierarchy API summaries.
    hierarchySummaries.forEach((summary, roomId) => {
      if (parentByRoomId.get(roomId) !== spaceItem.roomId) return;
      if (!isChildSpace(roomId, getRoom, spaceRooms, hierarchySummaries)) return;
      if (spaceItems.find((item) => item.roomId === roomId)) return;

      findAndCollectHierarchySpaces({
        roomId,
        content: { via: summary.via ?? [] },
        ts: 0,
        space: true,
        parentId: spaceItem.roomId,
      });
    });
  };
  findAndCollectHierarchySpaces(rootSpaceItem);

  spaceItems = [
    rootSpaceItem,
    ...spaceItems
      .filter((item) => item.roomId !== rootSpaceId)
      .sort(hierarchyItemTs)
      .sort(hierarchyItemByOrder),
  ];

  return spaceItems;
};

export type SpaceHierarchy = {
  space: HierarchyItemSpace;
  rooms?: HierarchyItemRoom[];
};
const getSpaceHierarchy = (
  rootSpaceId: string,
  spaceRooms: Set<string>,
  getRoom: (roomId: string) => Room | undefined,
  closedCategory: (spaceId: string) => boolean,
  hierarchySummaries: Map<string, IHierarchyRoom>
): SpaceHierarchy[] => {
  const spaceItems: HierarchyItemSpace[] = getHierarchySpaces(
    rootSpaceId,
    getRoom,
    spaceRooms,
    hierarchySummaries
  );
  const parentByRoomId = buildParentByRoomId(hierarchySummaries.values());

  const hierarchy: SpaceHierarchy[] = spaceItems.map((spaceItem) => {
    const space = getRoom(spaceItem.roomId);
    if (!space || closedCategory(spaceItem.roomId)) {
      if (!space && !closedCategory(spaceItem.roomId) && hierarchySummaries.size > 0) {
        const summaryRooms = getChildRoomsFromSummaries(
          spaceItem.roomId,
          hierarchySummaries,
          parentByRoomId,
          spaceRooms,
          getRoom
        );
        if (summaryRooms.length > 0) {
          return {
            space: spaceItem,
            rooms: summaryRooms.sort(hierarchyItemTs).sort(hierarchyItemByOrder),
          };
        }
      }
      return {
        space: spaceItem,
      };
    }
    const childEvents = getStateEvents(space, StateEvent.SpaceChild);
    const childItems: HierarchyItemRoom[] = [];
    childEvents.forEach((childEvent) => {
      if (!isValidChild(childEvent)) return;
      const childId = childEvent.getStateKey();
      if (!childId || !isRoomId(childId)) return;
      if (isChildSpace(childId, getRoom, spaceRooms, hierarchySummaries)) return;

      const childItem: HierarchyItemRoom = {
        roomId: childId,
        content: childEvent.getContent<MSpaceChildContent>(),
        ts: childEvent.getTs(),
        parentId: spaceItem.roomId,
      };
      childItems.push(childItem);
    });

    const summaryRooms = getChildRoomsFromSummaries(
      spaceItem.roomId,
      hierarchySummaries,
      parentByRoomId,
      spaceRooms,
      getRoom
    );
    const roomIds = new Set(childItems.map((i) => i.roomId));
    summaryRooms.forEach((item) => {
      if (!roomIds.has(item.roomId)) {
        childItems.push(item);
      }
    });

    return {
      space: spaceItem,
      rooms: childItems.sort(hierarchyItemTs).sort(hierarchyItemByOrder),
    };
  });

  return hierarchy;
};

export const useSpaceHierarchy = (
  spaceId: string,
  spaceRooms: Set<string>,
  getRoom: (roomId: string) => Room | undefined,
  closedCategory: (spaceId: string) => boolean,
  hierarchySummaries: Map<string, IHierarchyRoom>
): SpaceHierarchy[] => {
  const mx = useMatrixClient();
  const roomToParents = useAtomValue(roomToParentsAtom);

  const [hierarchyAtom] = useState(() =>
    atom(getSpaceHierarchy(spaceId, spaceRooms, getRoom, closedCategory, hierarchySummaries))
  );
  const [hierarchy, setHierarchy] = useAtom(hierarchyAtom);

  useEffect(() => {
    setHierarchy(
      getSpaceHierarchy(spaceId, spaceRooms, getRoom, closedCategory, hierarchySummaries)
    );
  }, [mx, spaceId, spaceRooms, setHierarchy, getRoom, closedCategory, hierarchySummaries]);

  useStateEventCallback(
    mx,
    useCallback(
      (mEvent) => {
        if (mEvent.getType() !== StateEvent.SpaceChild) return;
        const eventRoomId = mEvent.getRoomId();
        if (!eventRoomId) return;

        if (spaceId === eventRoomId || getAllParents(roomToParents, eventRoomId).has(spaceId)) {
          setHierarchy(
            getSpaceHierarchy(spaceId, spaceRooms, getRoom, closedCategory, hierarchySummaries)
          );
        }
      },
      [spaceId, roomToParents, setHierarchy, spaceRooms, getRoom, closedCategory, hierarchySummaries]
    )
  );

  return hierarchy;
};

const getSpaceJoinedHierarchy = (
  rootSpaceId: string,
  getRoom: GetRoomCallback,
  excludeRoom: (parentId: string, roomId: string) => boolean,
  sortRoomItems: (parentId: string, items: HierarchyItem[]) => HierarchyItem[]
): HierarchyItem[] => {
  const spaceItems: HierarchyItemSpace[] = getHierarchySpaces(
    rootSpaceId,
    getRoom,
    new Set(),
    new Map()
  );

  const hierarchy: HierarchyItem[] = spaceItems.flatMap((spaceItem) => {
    const space = getRoom(spaceItem.roomId);
    if (!space) {
      return [];
    }
    const joinedRoomEvents = getStateEvents(space, StateEvent.SpaceChild).filter((childEvent) => {
      if (!isValidChild(childEvent)) return false;
      const childId = childEvent.getStateKey();
      if (!childId || !isRoomId(childId)) return false;
      const room = getRoom(childId);
      if (!room || isSpace(room)) return false;

      return true;
    });

    if (joinedRoomEvents.length === 0) return [];

    const childItems: HierarchyItemRoom[] = [];
    joinedRoomEvents.forEach((childEvent) => {
      const childId = childEvent.getStateKey();
      if (!childId) return;

      if (excludeRoom(space.roomId, childId)) return;

      const childItem: HierarchyItemRoom = {
        roomId: childId,
        content: childEvent.getContent<MSpaceChildContent>(),
        ts: childEvent.getTs(),
        parentId: spaceItem.roomId,
      };
      childItems.push(childItem);
    });
    return [spaceItem, ...sortRoomItems(spaceItem.roomId, childItems)];
  });

  return hierarchy;
};

export const useSpaceJoinedHierarchy = (
  spaceId: string,
  getRoom: GetRoomCallback,
  excludeRoom: (parentId: string, roomId: string) => boolean,
  sortByActivity: (spaceId: string) => boolean
): HierarchyItem[] => {
  const mx = useMatrixClient();
  const roomToParents = useAtomValue(roomToParentsAtom);

  const sortRoomItems = useCallback(
    (sId: string, items: HierarchyItem[]) => {
      if (sortByActivity(sId)) {
        items.sort((a, b) => factoryRoomIdByActivity(mx)(a.roomId, b.roomId));
        return items;
      }
      items.sort(hierarchyItemTs).sort(hierarchyItemByOrder);
      return items;
    },
    [mx, sortByActivity]
  );

  const [hierarchyAtom] = useState(() =>
    atom(getSpaceJoinedHierarchy(spaceId, getRoom, excludeRoom, sortRoomItems))
  );
  const [hierarchy, setHierarchy] = useAtom(hierarchyAtom);

  useEffect(() => {
    setHierarchy(getSpaceJoinedHierarchy(spaceId, getRoom, excludeRoom, sortRoomItems));
  }, [mx, spaceId, setHierarchy, getRoom, excludeRoom, sortRoomItems]);

  useStateEventCallback(
    mx,
    useCallback(
      (mEvent) => {
        if (mEvent.getType() !== StateEvent.SpaceChild) return;
        const eventRoomId = mEvent.getRoomId();
        if (!eventRoomId) return;

        if (spaceId === eventRoomId || getAllParents(roomToParents, eventRoomId).has(spaceId)) {
          setHierarchy(getSpaceJoinedHierarchy(spaceId, getRoom, excludeRoom, sortRoomItems));
        }
      },
      [spaceId, roomToParents, setHierarchy, getRoom, excludeRoom, sortRoomItems]
    )
  );

  return hierarchy;
};

export type FetchSpaceHierarchyLevelData = {
  fetching: boolean;
  error: Error | null;
  rooms: Map<string, IHierarchyRoom>;
};
export const useFetchSpaceHierarchyLevel = (
  roomId: string,
  enable: boolean
): FetchSpaceHierarchyLevelData => {
  const mx = useMatrixClient();
  const pageNoRef = useRef(0);

  const fetchLevel: QueryFunction<
    Awaited<ReturnType<typeof mx.getRoomHierarchy>>,
    string[],
    string | undefined
  > = useCallback(
    ({ pageParam }) => mx.getRoomHierarchy(roomId, PER_PAGE_COUNT, 1, false, pageParam),
    [roomId, mx]
  );

  const queryResponse = useInfiniteQuery({
    refetchOnMount: enable,
    queryKey: [roomId, 'hierarchy_level'],
    initialPageParam: undefined,
    queryFn: fetchLevel,
    getNextPageParam: (result) => {
      if (result.next_batch) return result.next_batch;
      return undefined;
    },
    retry: 5,
    retryDelay: (failureCount, error) => {
      if (error instanceof MatrixError && error.errcode === ErrorCode.M_LIMIT_EXCEEDED) {
        const { retry_after_ms: delay } = error.data;
        if (typeof delay === 'number') {
          return delay;
        }
      }

      return 500 * failureCount;
    },
  });

  const { data, isLoading, isFetchingNextPage, error, fetchNextPage, hasNextPage } = queryResponse;

  useEffect(() => {
    if (
      hasNextPage &&
      pageNoRef.current <= MAX_AUTO_PAGE_COUNT &&
      !error &&
      data &&
      data.pages.length > 0
    ) {
      pageNoRef.current += 1;
      fetchNextPage();
    }
  }, [fetchNextPage, hasNextPage, data, error]);

  const rooms: Map<string, IHierarchyRoom> = useMemo(() => {
    const roomsMap: Map<string, IHierarchyRoom> = new Map();
    if (!data) return roomsMap;

    const rms = data.pages.flatMap((result) => result.rooms);
    rms.forEach((r) => {
      roomsMap.set(r.room_id, r);
    });

    return roomsMap;
  }, [data]);

  const fetching = isLoading || isFetchingNextPage;

  return {
    fetching,
    error,
    rooms,
  };
};
