keyboard/CustomKeyboard/KeyboardViewControllerHelp/KeyboardViewController+Suggestions.m

//
//  KeyboardViewController+Suggestions.m
//  CustomKeyboard
//
//  Created by Codex on 2026/02/22.
//

#import "KeyboardViewController+Private.h"

#import "KBBackspaceUndoManager.h"
#import "KBInputBufferManager.h"
#import "KBKeyBoardMainView.h"
#import "KBSuggestionEngine.h"

@implementation KeyboardViewController (Suggestions)

// MARK: - Suggestions

- (void)kb_updateCurrentWordWithInsertedText:(NSString *)text {
  if (text.length == 0) {
    return;
  }
  if ([self kb_isAlphabeticString:text]) {
    NSString *current = self.currentWord ?: @"";
    self.currentWord = [current stringByAppendingString:text];
    self.suppressSuggestions = NO;
    [self kb_updateSuggestionsForCurrentWord];
  } else {
    [self kb_clearCurrentWord];
  }
}

- (void)kb_clearCurrentWord {
  self.currentWord = @"";
  [self.keyBoardMainView kb_setSuggestions:@[]];
  self.suppressSuggestions = NO;
}

- (void)kb_scheduleContextRefreshResetSuppression:(BOOL)resetSuppression {
  dispatch_async(dispatch_get_main_queue(), ^{
    [self kb_refreshCurrentWordFromDocumentContextResetSuppression:
              resetSuppression];
  });
}

- (void)kb_refreshCurrentWordFromDocumentContextResetSuppression:
    (BOOL)resetSuppression {
  NSString *context = self.textDocumentProxy.documentContextBeforeInput ?: @"";
  NSString *word = [self kb_extractTrailingWordFromContext:context];
  self.currentWord = word ?: @"";
  if (resetSuppression) {
    self.suppressSuggestions = NO;
  }
  [self kb_updateSuggestionsForCurrentWord];
}

- (NSString *)kb_extractTrailingWordFromContext:(NSString *)context {
  if (context.length == 0) {
    return @"";
  }
  static NSCharacterSet *letters = nil;
  static dispatch_once_t onceToken;
  dispatch_once(&onceToken, ^{
    letters = [NSCharacterSet
        characterSetWithCharactersInString:
            @"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"];
  });

  NSInteger idx = (NSInteger)context.length - 1;
  while (idx >= 0) {
    unichar ch = [context characterAtIndex:(NSUInteger)idx];
    if (![letters characterIsMember:ch]) {
      break;
    }
    idx -= 1;
  }
  NSUInteger start = (NSUInteger)(idx + 1);
  if (start >= context.length) {
    return @"";
  }
  return [context substringFromIndex:start];
}

- (BOOL)kb_isAlphabeticString:(NSString *)text {
  if (text.length == 0) {
    return NO;
  }
  static NSCharacterSet *letters = nil;
  static dispatch_once_t onceToken;
  dispatch_once(&onceToken, ^{
    letters = [NSCharacterSet
        characterSetWithCharactersInString:
            @"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"];
  });
  for (NSUInteger i = 0; i < text.length; i++) {
    if (![letters characterIsMember:[text characterAtIndex:i]]) {
      return NO;
    }
  }
  return YES;
}

- (void)kb_updateSuggestionsForCurrentWord {
  NSString *prefix = self.currentWord ?: @"";
  if (prefix.length == 0) {
    [self.keyBoardMainView kb_setSuggestions:@[]];
    return;
  }
  if (self.suppressSuggestions) {
    [self.keyBoardMainView kb_setSuggestions:@[]];
    return;
  }
  NSArray<NSString *> *items =
      [self.suggestionEngine suggestionsForPrefix:prefix limit:5];
  NSArray<NSString *> *cased = [self kb_applyCaseToSuggestions:items
                                                        prefix:prefix];
  [self.keyBoardMainView kb_setSuggestions:cased];
}

- (NSArray<NSString *> *)kb_applyCaseToSuggestions:(NSArray<NSString *> *)items
                                            prefix:(NSString *)prefix {
  if (items.count == 0 || prefix.length == 0) {
    return items;
  }
  BOOL allUpper = [prefix isEqualToString:prefix.uppercaseString];
  BOOL firstUpper = [[prefix substringToIndex:1]
      isEqualToString:[[prefix substringToIndex:1] uppercaseString]];

  if (!allUpper && !firstUpper) {
    return items;
  }

  NSMutableArray<NSString *> *result =
      [NSMutableArray arrayWithCapacity:items.count];
  for (NSString *word in items) {
    if (allUpper) {
      [result addObject:word.uppercaseString];
    } else {
      NSString *first = [[word substringToIndex:1] uppercaseString];
      NSString *rest = (word.length > 1) ? [word substringFromIndex:1] : @"";
      [result addObject:[first stringByAppendingString:rest]];
    }
  }
  return result.copy;
}

// MARK: - KBKeyBoardMainViewDelegate (Suggestion)

- (void)keyBoardMainView:(KBKeyBoardMainView *)keyBoardMainView
     didSelectSuggestion:(NSString *)suggestion {
  if (suggestion.length == 0) {
    return;
  }
  NSDictionary *extra = @{@"suggestion_len" : @(suggestion.length)};
  //  [[KBMaiPointReporter sharedReporter]
  //      reportClickWithEventName:@"click_keyboard_suggestion_item"
  //                        pageId:@"keyboard_main_panel"
  //                     elementId:@"suggestion_item"
  //                         extra:extra
  //                    completion:nil];
  [[KBBackspaceUndoManager shared] registerNonClearAction];
  NSString *current = self.currentWord ?: @"";
  if (current.length > 0) {
    for (NSUInteger i = 0; i < current.length; i++) {
      [self.textDocumentProxy deleteBackward];
    }
  }
  [self.textDocumentProxy insertText:suggestion];
  self.currentWord = suggestion;
  [self.suggestionEngine recordSelection:suggestion];
  self.suppressSuggestions = YES;
  [self.keyBoardMainView kb_setSuggestions:@[]];
  [[KBInputBufferManager shared] replaceTailWithText:suggestion
                                         deleteCount:current.length];
}

@end