import { ITrackItemsMap, ITransition, ITrackItem } from "@designcombo/types";

type GroupElement = ITrackItem | ITransition;

export const groupTrackItems = (data: {
  trackItemIds: string[];
  transitionsMap: Record<string, ITransition>;
  trackItemsMap: ITrackItemsMap;
}): GroupElement[][] => {
  const { trackItemIds, transitionsMap, trackItemsMap } = data;

  // Create a map to track which items are part of transitions
  const itemTransitionMap = new Map<string, ITransition[]>();

  // Initialize transition maps
  Object.values(transitionsMap).forEach((transition) => {
    const { fromId, toId, kind } = transition;
    if (kind === "none") return; // Skip transitions of kind 'none'
    if (!itemTransitionMap.has(fromId)) itemTransitionMap.set(fromId, []);
    if (!itemTransitionMap.has(toId)) itemTransitionMap.set(toId, []);
    itemTransitionMap.get(fromId)?.push(transition);
    itemTransitionMap.get(toId)?.push(transition);
  });

  const groups: GroupElement[][] = [];
  const processed = new Set<string>();

  // Helper function to build a connected group starting from an item
  const buildGroup = (startItemId: string): GroupElement[] => {
    const group: GroupElement[] = [];
    let currentId = startItemId;

    while (currentId) {
      if (processed.has(currentId)) break;

      processed.add(currentId);
      const currentItem = trackItemsMap[currentId];
      group.push(currentItem);

      // Find transition from this item
      const transition = Object.values(transitionsMap).find(
        (t) => t.fromId === currentId && t.kind !== "none" // Filter here
      );
      if (!transition) break;

      group.push(transition);
      currentId = transition.toId;
    }

    return group;
  };

  // Process all items
  for (const itemId of trackItemIds) {
    if (processed.has(itemId)) continue;

    // If item is not part of any transition or is the start of a sequence
    if (
      !itemTransitionMap.has(itemId) ||
      !Object.values(transitionsMap).some((t) => t.toId === itemId)
    ) {
      const group = buildGroup(itemId);
      if (group.length > 0) {
        groups.push(group);
      }
    }
  }

  // Sort items within each group by display.from
  groups.forEach((group) => {
    group.sort((a, b) => {
      if ("display" in a && "display" in b) {
        return a.display.from - b.display.from;
      }
      return 0;
    });
  });

  return groups;
};