#import "RNSStackHeaderMenuCoordinator.h"
#import "RNSDefines.h"
#import "RNSStackHeaderIconResolver.h"

#import <React/RCTAssert.h>

@implementation RNSStackHeaderMenuCoordinator

+ (void)applyMenu:(RNSStackHeaderMenuData *)data
             toBarButtonItem:(UIBarButtonItem *)item
    withHeaderEventsDelegate:(id<RNSStackHeaderEventsDelegate>)delegate
                stateTracker:(RNSStackHeaderMenuToggleStateTracker *)tracker
             withImageLoader:(id<RNSImageLoading>)imageLoader
     menuInvalidatedCallback:(nullable void (^)(void))onMenuInvalidated
{
#if !TARGET_OS_TV || __TV_OS_VERSION_MAX_ALLOWED >= 170000
  if (@available(tvOS 17.0, *)) {
    item.menu = [self buildMenuFromData:data
                  withHeaderEventsDelegate:delegate
                              stateTracker:tracker
                       singleSelectionRoot:nil
        initialSingleSelectionStateClaimed:NULL
                           withImageLoader:imageLoader
                   menuInvalidatedCallback:onMenuInvalidated];
  }
#endif // !TARGET_OS_TV || __TV_OS_VERSION_MAX_ALLOWED >= 170000
}

+ (UIMenu *)buildMenuFromData:(RNSStackHeaderMenuData *)data
              withHeaderEventsDelegate:(id<RNSStackHeaderEventsDelegate>)delegate
                          stateTracker:(RNSStackHeaderMenuToggleStateTracker *)tracker
                   singleSelectionRoot:(nullable RNSStackHeaderMenuData *)singleSelectionRoot
    initialSingleSelectionStateClaimed:(BOOL *)initialSingleSelectionStateClaimed
                       withImageLoader:(id<RNSImageLoading>)imageLoader
               menuInvalidatedCallback:(nullable void (^)(void))onMenuInvalidated
{
  // Resolve singleSelection root: first menu in hierarchy with singleSelection becomes the root.
  // Only the root is set the singleSelection option - less things to check if sth goes wrong
  // and it shouldn't change the behavior
  // There can be at most one element with initialToggleState and we're checking that with
  // initialSingleSelectionStateClaimed
  UIMenuOptions options = 0;
  RNSStackHeaderMenuData *resolvedRoot = singleSelectionRoot;
  BOOL newInitialSingleSelectionStateClaimed = NO;
  if (resolvedRoot == nil && data.singleSelection) {
    options |= UIMenuOptionsSingleSelection;
    resolvedRoot = data;
    initialSingleSelectionStateClaimed = &newInitialSingleSelectionStateClaimed;
  }
  if (data.displayInline) {
    options |= UIMenuOptionsDisplayInline;
  }
#if RNS_IPHONE_OS_VERSION_AVAILABLE(17_0) || (TARGET_OS_TV && __TV_OS_VERSION_MAX_ALLOWED >= 170000)
  if (@available(iOS 17.0, tvOS 17.0, *)) {
    if (data.displayAsPalette) {
      options |= UIMenuOptionsDisplayAsPalette;
    }
  }
#endif

  NSMutableArray<UIMenuElement *> *elements = [NSMutableArray arrayWithCapacity:data.children.count];
  for (id<RNSStackHeaderMenuElement> child in data.children) {
    UIMenuElement *element = [self buildElementFromData:child
                               withHeaderEventsDelegate:delegate
                                           stateTracker:tracker
                                             parentMenu:data
                                    singleSelectionRoot:resolvedRoot
                     initialSingleSelectionStateClaimed:initialSingleSelectionStateClaimed
                                        withImageLoader:imageLoader
                                menuInvalidatedCallback:onMenuInvalidated];
    if (element != nil) {
      [elements addObject:element];
    }
  }

  UIImage *menuImage = nil;
  if (data.icon != nil) {
    UIImage *syncImage = [RNSStackHeaderIconResolver resolveIcon:data.icon
                                                 withImageLoader:imageLoader
                                            asyncCompletionBlock:^(UIImage *_Nullable image) {
                                              if (onMenuInvalidated) {
                                                onMenuInvalidated();
                                              }
                                            }];
    menuImage = syncImage ? syncImage : [RNSStackHeaderIconResolver placeholderImageForIcon:data.icon];
  }

  return [UIMenu menuWithTitle:data.title image:menuImage identifier:nil options:options children:elements];
}

+ (nullable UIMenuElement *)buildElementFromData:(id<RNSStackHeaderMenuElement>)element
                        withHeaderEventsDelegate:(id<RNSStackHeaderEventsDelegate>)delegate
                                    stateTracker:(RNSStackHeaderMenuToggleStateTracker *)tracker
                                      parentMenu:(RNSStackHeaderMenuData *)parentMenu
                             singleSelectionRoot:(nullable RNSStackHeaderMenuData *)singleSelectionRoot
              initialSingleSelectionStateClaimed:(BOOL *)initialSingleSelectionStateClaimed
                                 withImageLoader:(id<RNSImageLoading>)imageLoader
                         menuInvalidatedCallback:(nullable void (^)(void))onMenuInvalidated
{
  if ([element isKindOfClass:[RNSStackHeaderMenuData class]]) {
    return [self buildMenuFromData:(RNSStackHeaderMenuData *)element
                  withHeaderEventsDelegate:delegate
                              stateTracker:tracker
                       singleSelectionRoot:singleSelectionRoot
        initialSingleSelectionStateClaimed:initialSingleSelectionStateClaimed
                           withImageLoader:imageLoader
                   menuInvalidatedCallback:onMenuInvalidated];
  }

  if ([element isKindOfClass:[RNSStackHeaderMenuItemData class]]) {
    RNSStackHeaderMenuItemData *itemData = (RNSStackHeaderMenuItemData *)element;
    BOOL insideSingleSelection = singleSelectionRoot != nil;
    RNSMenuItemType effectiveItemType = [self resolveItemType:itemData.itemType
                                        insideSingleSelection:insideSingleSelection];

    RCTAssert(!(insideSingleSelection && effectiveItemType == RNSMenuItemTypeAction),
              @"[RNScreens] 'action' itemType is disallowed in singleSelection menus (id: %@)",
              itemData.menuElementId);

    if (effectiveItemType == RNSMenuItemTypeToggle) {
      if (insideSingleSelection && itemData.initialToggleState) {
        RCTAssert(
            !*initialSingleSelectionStateClaimed,
            @"[RNScreens] Multiple items with initialToggleState=YES in singleSelection menu hierarchy starting from \"%@\".",
            singleSelectionRoot.menuElementId);
        *initialSingleSelectionStateClaimed = YES;
      }
      return [self buildToggleFromData:itemData
                        withParentMenu:parentMenu
                   singleSelectionRoot:singleSelectionRoot
                    toggleStateTracker:tracker
                  headerEventsDelegate:delegate
                       withImageLoader:imageLoader
               menuInvalidatedCallback:onMenuInvalidated];
    }

    // it effective type is not 'toggle', then it is a regular action button that triggers onPress instead
    return [self buildActionFromData:itemData
            withHeaderEventsDelegate:delegate
                     withImageLoader:imageLoader
             menuInvalidatedCallback:onMenuInvalidated];
  }

  return nil;
}

/**
 Builds item which can be toggled on / off.

 Toggle works differently with singleSelection and without it. Normally, each toggle is separate and changing one
 triggers selection change event on direct parent, but when singleSelection is set somewhere in the parent *chain*,
 only one toggle in the whole hierarchy can be turned on at any point in time, and the single selection root,
 which doesn't need to be a direct parent, is notified.

 We don't send press event for toggle by design.
 */
+ (nullable UIMenuElement *)buildToggleFromData:(RNSStackHeaderMenuItemData *)data
                                 withParentMenu:(RNSStackHeaderMenuData *)parentMenu
                            singleSelectionRoot:(nullable RNSStackHeaderMenuData *)singleSelectionRoot
                             toggleStateTracker:(RNSStackHeaderMenuToggleStateTracker *)tracker
                           headerEventsDelegate:(id<RNSStackHeaderEventsDelegate>)delegate
                                withImageLoader:(id<RNSImageLoading>)imageLoader
                        menuInvalidatedCallback:(nullable void (^)(void))onMenuInvalidated
{
  BOOL isItemToggledOn = [tracker getToggleStateForItemWithId:data.menuElementId initialState:data.initialToggleState];
  BOOL insideSingleSelection = singleSelectionRoot != nil;

  NSString *eventMenuId = insideSingleSelection ? singleSelectionRoot.menuElementId : parentMenu.menuElementId;

  __weak id<RNSStackHeaderEventsDelegate> weakDelegate = delegate;

  UIImage *iconImage = nil;
  if (data.icon != nil) {
    UIImage *syncImage = [RNSStackHeaderIconResolver resolveIcon:data.icon
                                                 withImageLoader:imageLoader
                                            asyncCompletionBlock:^(UIImage *_Nullable image) {
                                              if (onMenuInvalidated) {
                                                onMenuInvalidated();
                                              }
                                            }];
    iconImage = syncImage ? syncImage : [RNSStackHeaderIconResolver placeholderImageForIcon:data.icon];
  }

  UIAction *toggleAction = [UIAction
      actionWithTitle:data.title
                image:iconImage
           identifier:nil
              handler:^(__kindof UIAction *_Nonnull action) {
                NSArray<NSString *> *toggleItemsIds = insideSingleSelection
                    ? [RNSStackHeaderMenuCoordinator getAllToggleItemsIdsInSingleSelectionHierarchy:singleSelectionRoot]
                    : [RNSStackHeaderMenuCoordinator getToggleItemsIdsInMenu:parentMenu];

                if (insideSingleSelection) {
                  if ([tracker toggleStateEquals:YES forItemWithId:data.menuElementId]) {
                    return;
                  }

                  [tracker selectItemWithId:data.menuElementId fromIds:toggleItemsIds];
                } else {
                  [tracker toggleItemWithId:data.menuElementId];
                }

                NSMutableArray<NSString *> *selectedIds = [NSMutableArray new];
                for (NSString *itemId in toggleItemsIds) {
                  if ([tracker getToggleStateForItemWithId:itemId initialState:NO]) {
                    [selectedIds addObject:itemId];
                  }
                }

                [weakDelegate didChangeSelectionForMenu:eventMenuId selectedMenuItemIds:selectedIds];

                if (onMenuInvalidated) {
                  onMenuInvalidated();
                }
              }];
  toggleAction.state = isItemToggledOn ? UIMenuElementStateOn : UIMenuElementStateOff;

  [self decorateActionKeepsMenuPresented:toggleAction withData:data];

  return toggleAction;
}

+ (nullable UIMenuElement *)buildActionFromData:(RNSStackHeaderMenuItemData *)data
                       withHeaderEventsDelegate:(id<RNSStackHeaderEventsDelegate>)delegate
                                withImageLoader:(id<RNSImageLoading>)imageLoader
                        menuInvalidatedCallback:(nullable void (^)(void))onMenuInvalidated
{
  __weak id<RNSStackHeaderEventsDelegate> weakDelegate = delegate;

  UIImage *iconImage = nil;
  if (data.icon != nil) {
    UIImage *syncImage = [RNSStackHeaderIconResolver resolveIcon:data.icon
                                                 withImageLoader:imageLoader
                                            asyncCompletionBlock:^(UIImage *_Nullable image) {
                                              if (onMenuInvalidated) {
                                                onMenuInvalidated();
                                              }
                                            }];
    iconImage = syncImage ? syncImage : [RNSStackHeaderIconResolver placeholderImageForIcon:data.icon];
  }

  UIAction *action = [UIAction actionWithTitle:data.title
                                         image:iconImage
                                    identifier:nil
                                       handler:^(__kindof UIAction *_Nonnull action) {
                                         [weakDelegate didPressMenuItem:data.menuElementId];
                                       }];

  [self decorateActionKeepsMenuPresented:action withData:data];

  return action;
}

#pragma mark - Helpers

+ (RNSMenuItemType)resolveItemType:(RNSMenuItemType)itemType insideSingleSelection:(BOOL)insideSingleSelection
{
  if (itemType == RNSMenuItemTypeAutomatic) {
    return insideSingleSelection ? RNSMenuItemTypeToggle : RNSMenuItemTypeAction;
  }
  return itemType;
}

+ (NSArray<NSString *> *)getToggleItemsIdsInMenu:(RNSStackHeaderMenuData *)menu
{
  NSMutableArray<NSString *> *ids = [NSMutableArray new];
  [self collectToggleItemsIdsFromMenu:menu intoArray:ids insideSingleSelection:NO];
  return ids;
}

+ (NSArray<NSString *> *)getAllToggleItemsIdsInSingleSelectionHierarchy:(RNSStackHeaderMenuData *)menu
{
  NSMutableArray<NSString *> *ids = [NSMutableArray new];
  [self collectToggleItemsIdsFromMenu:menu intoArray:ids insideSingleSelection:YES];
  return ids;
}

+ (void)collectToggleItemsIdsFromMenu:(RNSStackHeaderMenuData *)menu
                            intoArray:(NSMutableArray<NSString *> *)ids
                insideSingleSelection:(bool)insideSingleSelection
{
  for (id<RNSStackHeaderMenuElement> child in menu.children) {
    if ([child isKindOfClass:[RNSStackHeaderMenuItemData class]]) {
      RNSStackHeaderMenuItemData *item = (RNSStackHeaderMenuItemData *)child;
      RNSMenuItemType effectiveItemType = [self resolveItemType:item.itemType
                                          insideSingleSelection:insideSingleSelection];
      if (effectiveItemType == RNSMenuItemTypeToggle) {
        [ids addObject:item.menuElementId];
      }
    } else if (insideSingleSelection && [child isKindOfClass:[RNSStackHeaderMenuData class]]) {
      [self collectToggleItemsIdsFromMenu:(RNSStackHeaderMenuData *)child
                                intoArray:ids
                    insideSingleSelection:insideSingleSelection];
    }
  }
}

+ (void)decorateActionKeepsMenuPresented:(UIAction *)action withData:(RNSStackHeaderMenuItemData *)data
{
  if (data.keepsMenuPresented) {
#if RNS_IPHONE_OS_VERSION_AVAILABLE(16_0) || (TARGET_OS_TV && __TV_OS_VERSION_MAX_ALLOWED >= 160000)
    if (@available(iOS 16.0, tvOS 16.0, *)) {
      action.attributes |= UIMenuElementAttributesKeepsMenuPresented;
    }
#endif
  }
}

+ (RNSStackHeaderMenuData *)menu:(RNSStackHeaderMenuData *)menu
            replacingChildWithId:(NSString *)elementId
                     withElement:(id<RNSStackHeaderMenuElement>)newElement
{
  BOOL changed = NO;
  NSMutableArray<id<RNSStackHeaderMenuElement>> *newChildren = [NSMutableArray arrayWithCapacity:menu.children.count];

  for (id<RNSStackHeaderMenuElement> child in menu.children) {
    if ([child.menuElementId isEqualToString:elementId]) {
      [newChildren addObject:newElement];
      changed = YES;
    } else if ([child isKindOfClass:[RNSStackHeaderMenuData class]]) {
      RNSStackHeaderMenuData *rebuilt = [self menu:(RNSStackHeaderMenuData *)child
                              replacingChildWithId:elementId
                                       withElement:newElement];
      [newChildren addObject:rebuilt];
      if (rebuilt != child) {
        changed = YES;
      }
    } else {
      [newChildren addObject:child];
    }
  }

  if (!changed) {
    return menu;
  }

  return [[RNSStackHeaderMenuData alloc] initWithId:menu.menuElementId
                                              title:menu.title
                                    singleSelection:menu.singleSelection
                                      displayInline:menu.displayInline
                                   displayAsPalette:menu.displayAsPalette
                                           children:newChildren
                                               icon:menu.icon];
}

@end