keyboard/CustomKeyboard/KeyboardViewControllerHelp/KeyboardViewController+Chat.m

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

#import "KeyboardViewController+Private.h"

#import "KBChatLimitPopView.h"
#import "KBChatMessage.h"
#import "KBChatPanelView.h"
#import "KBFullAccessManager.h"
#import "KBHostAppLauncher.h"
#import "KBInputBufferManager.h"
#import "KBNetworkManager.h"
#import "KBVM.h"
#import "Masonry.h"

#import <AVFoundation/AVFoundation.h>

static const NSUInteger kKBChatMessageLimit = 6;

@implementation KeyboardViewController (Chat)

#pragma mark - KBChatPanelViewDelegate

- (void)chatPanelView:(KBChatPanelView *)view didSendText:(NSString *)text {
  NSString *trim =
      [text stringByTrimmingCharactersInSet:
                [NSCharacterSet whitespaceAndNewlineCharacterSet]];
  if (trim.length == 0) {
    return;
  }
  [self kb_sendChatText:trim];
}

- (void)chatPanelView:(KBChatPanelView *)view
       didTapMessage:(KBChatMessage *)message {
  if (message.audioFilePath.length == 0) {
    return;
  }
  [self kb_playChatAudioAtPath:message.audioFilePath];
}

- (void)chatPanelView:(KBChatPanelView *)view
    didTapVoiceButtonForMessage:(KBChatMessage *)message {
  if (!message)
    return;

  // 如果有 audioData，直接播放
  if (message.audioData && message.audioData.length > 0) {
    [self kb_playChatAudioData:message.audioData];
    return;
  }

  // 如果有 audioFilePath，播放文件
  if (message.audioFilePath.length > 0) {
    [self kb_playChatAudioAtPath:message.audioFilePath];
    return;
  }

  NSLog(@"[Keyboard] 没有音频数据可播放");
}

- (void)chatPanelViewDidTapClose:(KBChatPanelView *)view {
  // 清空 chatPanelView 内部的消息
  [view kb_reloadWithMessages:@[]];
  if (self.chatAudioPlayer.isPlaying) {
    [self.chatAudioPlayer stop];
  }
  self.chatAudioPlayer = nil;
  [self kb_setPanelMode:KBKeyboardPanelModeMain animated:YES];
}

#pragma mark - Chat Helpers

- (void)kb_handleChatSendAction {
  if (!self.chatPanelVisible) {
    return;
  }
  [[KBInputBufferManager shared]
      refreshFromProxyIfPossible:self.textDocumentProxy];
  NSString *fullText = [KBInputBufferManager shared].liveText ?: @"";

  // 去掉打开聊天面板前宿主输入框里已有的基线文本，只取新增部分
  NSString *baseline = self.chatPanelBaselineText ?: @"";
  NSString *rawText = fullText;
  if (baseline.length > 0 && [fullText hasPrefix:baseline]) {
    rawText = [fullText substringFromIndex:baseline.length];
  }

  NSString *trim =
      [rawText stringByTrimmingCharactersInSet:
                  [NSCharacterSet whitespaceAndNewlineCharacterSet]];
  NSString *textToClear = rawText;
  if (trim.length == 0) {
    // 兼容「先输入再打开聊天面板」场景：
    // 此时新增文本为空，但当前输入框已有可发送内容，应该允许直接发送。
    NSString *fullTrim =
        [fullText stringByTrimmingCharactersInSet:
                      [NSCharacterSet whitespaceAndNewlineCharacterSet]];
    if (fullTrim.length > 0) {
      trim = fullTrim;
      textToClear = fullText;
    }
  }
  if (trim.length == 0) {
    [KBHUD showInfo:KBLocalized(@"请输入内容")];
    return;
  }
  [self kb_sendChatText:trim];
  // 默认只清新增文本；若命中兜底则清当前全文，避免“已发送但输入框残留”。
  [self kb_clearHostInputForText:textToClear];
}

- (void)kb_sendChatText:(NSString *)text {
  if (text.length == 0) {
    return;
  }
  NSLog(@"[KB] 发送消息: %@", text);

  KBChatMessage *outgoing = [KBChatMessage userMessageWithText:text];
  outgoing.avatarURL = [self kb_sharedUserAvatarURL];
  [self.chatPanelView kb_addUserMessage:text];
  [self kb_prefetchAvatarForMessage:outgoing];

  if (![[KBFullAccessManager shared] ensureFullAccessOrGuideInView:self.view]) {
    [KBHUD showInfo:KBLocalized(@"请开启完全访问后使用")];
    return;
  }

  // 添加 loading 消息
  [self.chatPanelView kb_addLoadingAssistantMessage];

  // 调用新的聊天接口
  [self kb_requestChatMessageWithContent:text];
}

#pragma mark - Chat Limit Pop

- (void)kb_showChatLimitPopWithMessage:(NSString *)message {
  [self kb_dismissChatLimitPop];

  UIControl *mask = [[UIControl alloc] init];
  mask.backgroundColor = [[UIColor blackColor] colorWithAlphaComponent:0.4];
  mask.alpha = 0.0;
  [mask addTarget:self
                action:@selector(kb_dismissChatLimitPop)
      forControlEvents:UIControlEventTouchUpInside];
  [self.contentView addSubview:mask];
  [mask mas_makeConstraints:^(MASConstraintMaker *make) {
    make.edges.equalTo(self.contentView);
  }];

  CGFloat width = 252.0;
  CGFloat height = 252.0 + 18.0 + 53.0 + 18.0 + 28.0;
  KBChatLimitPopView *content =
      [[KBChatLimitPopView alloc] initWithFrame:CGRectMake(0, 0, width, height)];
  content.message = message ?: @"";
  content.delegate = self;
  [mask addSubview:content];
  [content mas_makeConstraints:^(MASConstraintMaker *make) {
    make.center.equalTo(mask);
    make.width.mas_equalTo(width);
    make.height.mas_equalTo(height);
  }];

  self.chatLimitMaskView = mask;
  [self.contentView bringSubviewToFront:mask];
  [UIView animateWithDuration:0.18
                   animations:^{
                     mask.alpha = 1.0;
                   }];
}

- (void)kb_dismissChatLimitPop {
  if (!self.chatLimitMaskView) {
    return;
  }
  UIControl *mask = self.chatLimitMaskView;
  self.chatLimitMaskView = nil;
  [UIView animateWithDuration:0.15
                   animations:^{
                     mask.alpha = 0.0;
                   }
                   completion:^(__unused BOOL finished) {
                     [mask removeFromSuperview];
                   }];
}

- (void)kb_clearHostInputForText:(NSString *)text {
  if (text.length == 0) {
    return;
  }
  NSUInteger count = [self kb_composedCharacterCountForString:text];
  for (NSUInteger i = 0; i < count; i++) {
    [self.textDocumentProxy deleteBackward];
  }
  [[KBInputBufferManager shared] clearAllLiveText];
  [self kb_clearCurrentWord];
}

- (NSUInteger)kb_composedCharacterCountForString:(NSString *)text {
  if (text.length == 0) {
    return 0;
  }
  __block NSUInteger count = 0;
  [text enumerateSubstringsInRange:NSMakeRange(0, text.length)
                           options:NSStringEnumerationByComposedCharacterSequences
                        usingBlock:^(__unused NSString *substring,
                                     __unused NSRange substringRange,
                                     __unused NSRange enclosingRange,
                                     __unused BOOL *stop) {
                          count += 1;
                        }];
  return count;
}

- (NSString *)kb_sharedUserAvatarURL {
  NSUserDefaults *ud = [[NSUserDefaults alloc] initWithSuiteName:AppGroup];
  NSString *url = [ud stringForKey:AppGroup_UserAvatarURL];
  return url ?: @"";
}

- (void)kb_prefetchAvatarForMessage:(KBChatMessage *)message {
  if (!message || message.avatarImage) {
    return;
  }
  NSString *urlString = message.avatarURL ?: @"";
  if (urlString.length == 0) {
    return;
  }
  if (![[KBFullAccessManager shared] hasFullAccess]) {
    return;
  }
  __weak typeof(self) weakSelf = self;
  [[KBVM shared] downloadAvatarFromURL:urlString
                             completion:^(UIImage *image, NSError *error) {
                               __strong typeof(weakSelf) self = weakSelf;
                               if (!self || !image)
                                 return;

                               message.avatarImage = image;
                               [self kb_reloadChatRowForMessage:message];
                             }];
}

- (void)kb_reloadChatRowForMessage:(KBChatMessage *)message {
  // 头像预加载完成后不需要刷新表格
  // 因为键盘扩展的聊天面板不显示头像，所以这里直接返回
  // 如果将来需要显示头像，可以只刷新特定行而不是整个表格
}

- (void)kb_requestChatAudioForText:(NSString *)text {
  NSString *mockPath = [self kb_mockChatAudioPath];
  if (mockPath.length > 0) {
    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.35 * NSEC_PER_SEC)),
                   dispatch_get_main_queue(), ^{
                     NSString *displayText = KBLocalized(@"语音回复");
                     KBChatMessage *incoming =
                         [KBChatMessage messageWithText:displayText
                                                outgoing:NO
                                           audioFilePath:mockPath];
                     incoming.displayName = KBLocalized(@"AI助手");
                     [self kb_appendChatMessage:incoming];
                     [self kb_playChatAudioAtPath:mockPath];
                   });
    return;
  }
  NSDictionary *payload = @{@"message" : text ?: @""};
  __weak typeof(self) weakSelf = self;
  [[KBNetworkManager shared] POST:API_AI_TALK
                         jsonBody:payload
                          headers:nil
                       completion:^(NSDictionary *json, NSURLResponse *response,
                                    NSError *error) {
                         dispatch_async(dispatch_get_main_queue(), ^{
                           __strong typeof(weakSelf) self = weakSelf;
                           if (!self) {
                             return;
                           }
                           if (error) {
                             NSString *tip = error.localizedDescription
                                                 ?: KBLocalized(@"请求失败");
                             [KBHUD showInfo:tip];
                             return;
                           }
                           NSString *displayText =
                               [self kb_chatTextFromJSON:json];
                           NSString *audioURL =
                               [self kb_chatAudioURLFromJSON:json];
                           NSString *audioBase64 =
                               [self kb_chatAudioBase64FromJSON:json];
                           if (audioURL.length > 0) {
                             [self kb_downloadChatAudioFromURL:audioURL
                                                   displayText:displayText];
                             return;
                           }
                           if (audioBase64.length > 0) {
                             NSData *data = [[NSData alloc]
                                 initWithBase64EncodedString:audioBase64
                                                      options:0];
                             if (data.length == 0) {
                               [KBHUD showInfo:KBLocalized(@"音频数据解析失败")];
                               return;
                             }
                             [self kb_handleChatAudioData:data
                                              fileExtension:@"m4a"
                                                 displayText:displayText];
                             return;
                           }
                           [KBHUD showInfo:KBLocalized(@"未获取到音频文件")];
                         });
                       }];
}

#pragma mark - New Chat API (with typewriter effect and audio preload)

/// 调用新的聊天接口（返回文本和 audioId）
- (void)kb_requestChatMessageWithContent:(NSString *)content {
  if (content.length == 0) {
    [self.chatPanelView kb_removeLoadingAssistantMessage];
    return;
  }

  NSInteger companionId = [[KBVM shared] selectedCompanionIdFromAppGroup];
  NSLog(@"[KB] 请求聊天: companionId=%ld", (long)companionId);

  __weak typeof(self) weakSelf = self;
  [[KBVM shared] sendChatMessageWithContent:content
                                companionId:companionId
                                 completion:^(KBChatResponse *response) {
                                   __strong typeof(weakSelf) self = weakSelf;
                                   if (!self)
                                     return;

                                   if (response.code != 0) {
                                     if (response.code == 50030) {
                                       NSLog(@"[KB] ⚠️ 次数用尽: %@",
                                             response.message);
                                       [self.chatPanelView
                                           kb_removeLoadingAssistantMessage];
                                       [self kb_showChatLimitPopWithMessage:
                                                 response.message];
                                       return;
                                     }
                                     NSLog(@"[KB] ❌ 请求失败: %@",
                                           response.message);
                                     [self.chatPanelView
                                         kb_removeLoadingAssistantMessage];
                                     [KBHUD showInfo:response.message
                                                        ?: KBLocalized(@"请求失败")];
                                     return;
                                   }

                                   NSLog(@"[KB] ✅ 收到回复: %@",
                                         response.data.aiResponse);

                                   if (response.data.aiResponse.length == 0) {
                                     [self.chatPanelView
                                         kb_removeLoadingAssistantMessage];
                                     [KBHUD showInfo:KBLocalized(@"未获取到回复内容")];
                                     return;
                                   }

                                   // 添加 AI 消息（带打字机效果）
                                   NSLog(@"[KB] 准备添加 AI 消息");
                                   [self.chatPanelView
                                       kb_addAssistantMessage:response.data.aiResponse
                                                     audioId:response.data.audioId];
                                   NSLog(@"[KB] AI 消息添加完成");

                                   // 如果有 audioId，开始预加载音频
                                   if (response.data.audioId.length > 0) {
                                     [self kb_preloadAudioWithAudioId:
                                               response.data.audioId];
                                   }
                                 }];
}

/// 从 AppGroup 获取选中的 persona companionId
- (NSInteger)kb_selectedCompanionId {
  return [[KBVM shared] selectedCompanionIdFromAppGroup];
}

#pragma mark - Audio Preload

/// 预加载音频（轮询获取 audioURL）
- (void)kb_preloadAudioWithAudioId:(NSString *)audioId {
  if (audioId.length == 0)
    return;

  NSLog(@"[Keyboard] 开始预加载音频，audioId: %@", audioId);

  __weak typeof(self) weakSelf = self;
  [[KBVM shared] pollAudioURLWithAudioId:audioId
                              maxRetries:10
                                interval:1.0
                              completion:^(KBAudioResponse *response) {
                                __strong typeof(weakSelf) self = weakSelf;
                                if (!self)
                                  return;

                                if (!response.success ||
                                    response.audioURL.length == 0) {
                                  NSLog(@"[Keyboard] ❌ 预加载音频 URL 获取失败: %@",
                                        response.errorMessage);
                                  return;
                                }

                                NSLog(@"[Keyboard] ✅ 预加载音频 URL 获取成功");

                                // 下载音频
                                [[KBVM shared]
                                    downloadAudioFromURL:response.audioURL
                                             completion:^(
                                                 KBAudioResponse *audioResponse) {
                                               if (!audioResponse.success) {
                                                 NSLog(@"[Keyboard] ❌ 预加载音频下载失败: %@",
                                                       audioResponse.errorMessage);
                                                 return;
                                               }

                                               // 更新最后一条 AI 消息的音频数据
                                               [self.chatPanelView
                                                   kb_updateLastAssistantMessageWithAudioData:
                                                       audioResponse.audioData
                                                                                 duration:
                                                                                     audioResponse.duration];
                                               NSLog(@"[Keyboard] ✅ 预加载音频完成，音频时长: %.2f秒",
                                                     audioResponse.duration);
                                             }];
                              }];
}

- (void)kb_downloadChatAudioFromURL:(NSString *)audioURL
                        displayText:(NSString *)displayText {
  __weak typeof(self) weakSelf = self;
  [[KBVM shared] downloadAudioFromURL:audioURL
                           completion:^(KBAudioResponse *response) {
                             __strong typeof(weakSelf) self = weakSelf;
                             if (!self)
                               return;

                             if (!response.success) {
                               [KBHUD showInfo:response.errorMessage
                                                   ?: KBLocalized(@"下载失败")];
                               return;
                             }

                             if (!response.audioData ||
                                 response.audioData.length == 0) {
                               [KBHUD showInfo:KBLocalized(@"未获取到音频数据")];
                               return;
                             }

                             NSString *ext = @"m4a";
                             NSURL *url = [NSURL URLWithString:audioURL];
                             if (url.pathExtension.length > 0) {
                               ext = url.pathExtension;
                             }
                             [self kb_handleChatAudioData:response.audioData
                                            fileExtension:ext
                                               displayText:displayText];
                           }];
}

- (void)kb_handleChatAudioData:(NSData *)data
                 fileExtension:(NSString *)extension
                    displayText:(NSString *)displayText {
  if (data.length == 0) {
    [KBHUD showInfo:KBLocalized(@"音频数据为空")];
    return;
  }
  NSString *ext = extension.length > 0 ? extension : @"m4a";
  NSString *fileName = [NSString
      stringWithFormat:@"kb_chat_%@.%@",
                       @((long long)([NSDate date].timeIntervalSince1970 *
                                     1000)),
                       ext];
  NSString *filePath =
      [NSTemporaryDirectory() stringByAppendingPathComponent:fileName];
  if (![data writeToFile:filePath atomically:YES]) {
    [KBHUD showInfo:KBLocalized(@"音频保存失败")];
    return;
  }
  NSString *text =
      displayText.length > 0 ? displayText : KBLocalized(@"语音消息");
  KBChatMessage *incoming =
      [KBChatMessage messageWithText:text outgoing:NO audioFilePath:filePath];
  incoming.displayName = KBLocalized(@"AI助手");
  [self kb_appendChatMessage:incoming];
}

- (void)kb_appendChatMessage:(KBChatMessage *)message {
  if (!message) {
    return;
  }
  [self.chatMessages addObject:message];
  if (self.chatMessages.count > kKBChatMessageLimit) {
    NSUInteger overflow = self.chatMessages.count - kKBChatMessageLimit;
    NSArray<KBChatMessage *> *removed =
        [self.chatMessages subarrayWithRange:NSMakeRange(0, overflow)];
    [self.chatMessages removeObjectsInRange:NSMakeRange(0, overflow)];
    for (KBChatMessage *msg in removed) {
      if (msg.audioFilePath.length > 0) {
        NSString *tmpRoot = NSTemporaryDirectory();
        if (tmpRoot.length > 0 && [msg.audioFilePath hasPrefix:tmpRoot]) {
          [[NSFileManager defaultManager] removeItemAtPath:msg.audioFilePath
                                                     error:nil];
        }
      }
    }
  }
  [self.chatPanelView kb_reloadWithMessages:self.chatMessages];
}

- (NSString *)kb_mockChatAudioPath {
  NSString *path = [[NSBundle mainBundle] pathForResource:@"ai_test"
                                                   ofType:@"m4a"];
  return path ?: @"";
}

- (NSString *)kb_chatTextFromJSON:(NSDictionary *)json {
  NSDictionary *data = [self kb_chatDataDictionaryFromJSON:json];
  NSString *text =
      [self kb_stringValueInDict:data keys:@[ @"text", @"message", @"content" ]];
  if (text.length == 0) {
    text = [self kb_stringValueInDict:json
                                  keys:@[ @"text", @"message", @"content" ]];
  }
  return text ?: @"";
}

- (NSString *)kb_chatAudioURLFromJSON:(NSDictionary *)json {
  NSDictionary *data = [self kb_chatDataDictionaryFromJSON:json];
  NSArray<NSString *> *keys =
      @[ @"audioUrl", @"audioURL", @"audio_url", @"url", @"fileUrl",
         @"file_url", @"audioFileUrl", @"audio_file_url" ];
  NSString *url = [self kb_stringValueInDict:data keys:keys];
  if (url.length == 0) {
    url = [self kb_stringValueInDict:json keys:keys];
  }
  return url ?: @"";
}

- (NSString *)kb_chatAudioBase64FromJSON:(NSDictionary *)json {
  NSDictionary *data = [self kb_chatDataDictionaryFromJSON:json];
  NSArray<NSString *> *keys =
      @[ @"audioBase64", @"audio_base64", @"audioData", @"audio_data",
         @"base64" ];
  NSString *b64 = [self kb_stringValueInDict:data keys:keys];
  if (b64.length == 0) {
    b64 = [self kb_stringValueInDict:json keys:keys];
  }
  return b64 ?: @"";
}

- (NSDictionary *)kb_chatDataDictionaryFromJSON:(NSDictionary *)json {
  if (![json isKindOfClass:[NSDictionary class]]) {
    return @{};
  }
  id dataObj = json[@"data"] ?: json[@"result"] ?: json[@"response"];
  if ([dataObj isKindOfClass:[NSDictionary class]]) {
    return (NSDictionary *)dataObj;
  }
  return @{};
}

- (NSString *)kb_stringValueInDict:(NSDictionary *)dict
                               keys:(NSArray<NSString *> *)keys {
  if (![dict isKindOfClass:[NSDictionary class]]) {
    return @"";
  }
  for (NSString *key in keys) {
    id value = dict[key];
    if ([value isKindOfClass:[NSString class]] &&
        ((NSString *)value).length > 0) {
      return (NSString *)value;
    }
  }
  return @"";
}

- (void)kb_playChatAudioAtPath:(NSString *)path {
  if (path.length == 0) {
    return;
  }
  NSURL *url = [NSURL fileURLWithPath:path];
  if (![NSFileManager.defaultManager fileExistsAtPath:path]) {
    [KBHUD showInfo:KBLocalized(@"音频文件不存在")];
    return;
  }

  if (self.chatAudioPlayer && self.chatAudioPlayer.isPlaying) {
    NSURL *currentURL = self.chatAudioPlayer.url;
    if ([currentURL isEqual:url]) {
      [self.chatAudioPlayer stop];
      self.chatAudioPlayer = nil;
      return;
    }
    [self.chatAudioPlayer stop];
    self.chatAudioPlayer = nil;
  }

  NSError *sessionError = nil;
  AVAudioSession *session = [AVAudioSession sharedInstance];
  if ([session respondsToSelector:@selector(setCategory:options:error:)]) {
    [session setCategory:AVAudioSessionCategoryPlayback
             withOptions:AVAudioSessionCategoryOptionDuckOthers
                   error:&sessionError];
  } else {
    [session setCategory:AVAudioSessionCategoryPlayback error:&sessionError];
  }
  [session setActive:YES error:nil];

  NSError *playerError = nil;
  AVAudioPlayer *player =
      [[AVAudioPlayer alloc] initWithContentsOfURL:url error:&playerError];
  if (playerError || !player) {
    [KBHUD showInfo:KBLocalized(@"音频播放失败")];
    return;
  }
  self.chatAudioPlayer = player;
  [player prepareToPlay];
  [player play];
}

/// 播放音频数据
- (void)kb_playChatAudioData:(NSData *)audioData {
  if (!audioData || audioData.length == 0) {
    NSLog(@"[Keyboard] 音频数据为空");
    return;
  }

  // 如果正在播放，先停止
  if (self.chatAudioPlayer && self.chatAudioPlayer.isPlaying) {
    [self.chatAudioPlayer stop];
    self.chatAudioPlayer = nil;
  }

  // 配置音频会话
  NSError *sessionError = nil;
  AVAudioSession *session = [AVAudioSession sharedInstance];
  if ([session respondsToSelector:@selector(setCategory:options:error:)]) {
    [session setCategory:AVAudioSessionCategoryPlayback
             withOptions:AVAudioSessionCategoryOptionDuckOthers
                   error:&sessionError];
  } else {
    [session setCategory:AVAudioSessionCategoryPlayback error:&sessionError];
  }
  [session setActive:YES error:nil];

  // 创建播放器
  NSError *playerError = nil;
  AVAudioPlayer *player =
      [[AVAudioPlayer alloc] initWithData:audioData error:&playerError];
  if (playerError || !player) {
    NSLog(@"[Keyboard] 音频播放器初始化失败: %@",
          playerError.localizedDescription);
    [KBHUD showInfo:KBLocalized(@"音频播放失败")];
    return;
  }

  self.chatAudioPlayer = player;
  player.volume = 1.0;
  [player prepareToPlay];
  [player play];

  NSLog(@"[Keyboard] 开始播放音频，时长: %.2f秒", player.duration);
}

#pragma mark - KBChatLimitPopViewDelegate

- (void)chatLimitPopViewDidTapCancel:(KBChatLimitPopView *)view {
  [self kb_dismissChatLimitPop];
}

- (void)chatLimitPopViewDidTapRecharge:(KBChatLimitPopView *)view {
  [self kb_dismissChatLimitPop];
  NSString *urlString =
      [NSString stringWithFormat:@"%@://recharge?src=keyboard&vipType=svip",
                                 KB_APP_SCHEME];
  NSURL *scheme = [NSURL URLWithString:urlString];
  BOOL success = [KBHostAppLauncher openHostAppURL:scheme
                                     fromResponder:self.view];
  if (!success) {
    [KBHUD showInfo:KBLocalized(@"Please open the App to finish purchase")];
  }
}

@end