keyboard/keyBoard/Class/AiTalk/V/KBChatTableView.m

//
//  KBChatTableView.m
//  keyBoard
//
//  Created by Kiro on 2026/1/23.
//

#import "KBChatTableView.h"
#import "KBAiChatMessage.h"
#import "KBChatUserMessageCell.h"
#import "KBChatAssistantMessageCell.h"
#import "KBChatTimeCell.h"
#import "AiVM.h"
#import <Masonry/Masonry.h>
#import <AVFoundation/AVFoundation.h>

static NSString * const kUserCellIdentifier = @"KBChatUserMessageCell";
static NSString * const kAssistantCellIdentifier = @"KBChatAssistantMessageCell";
static NSString * const kTimeCellIdentifier = @"KBChatTimeCell";

/// 时间戳显示间隔（秒）
static const NSTimeInterval kTimestampInterval = 5 * 60; // 5 分钟

@interface KBChatTableView () <UITableViewDataSource, UITableViewDelegate, KBChatAssistantMessageCellDelegate, AVAudioPlayerDelegate>

@property (nonatomic, strong) UITableView *tableView;
@property (nonatomic, strong) NSMutableArray<KBAiChatMessage *> *messages;
@property (nonatomic, strong) AVAudioPlayer *audioPlayer;
@property (nonatomic, strong) NSIndexPath *playingCellIndexPath;
@property (nonatomic, strong) AiVM *aiVM;

@end

@implementation KBChatTableView

- (instancetype)initWithFrame:(CGRect)frame {
    self = [super initWithFrame:frame];
    if (self) {
        [self setup];
    }
    return self;
}

- (instancetype)initWithCoder:(NSCoder *)coder {
    self = [super initWithCoder:coder];
    if (self) {
        [self setup];
    }
    return self;
}

- (void)setup {
    self.messages = [[NSMutableArray alloc] init];
    self.aiVM = [[AiVM alloc] init];
    
    // 创建 TableView
    self.tableView = [[UITableView alloc] initWithFrame:self.bounds
                                                  style:UITableViewStylePlain];
    self.tableView.dataSource = self;
    self.tableView.delegate = self;
    self.tableView.separatorStyle = UITableViewCellSeparatorStyleNone;
    self.tableView.backgroundColor = [UIColor clearColor];
    self.tableView.estimatedRowHeight = 60;
    self.tableView.rowHeight = UITableViewAutomaticDimension;
    self.tableView.showsVerticalScrollIndicator = YES;
    [self addSubview:self.tableView];
    
    // 注册 Cell
    [self.tableView registerClass:[KBChatUserMessageCell class]
           forCellReuseIdentifier:kUserCellIdentifier];
    [self.tableView registerClass:[KBChatAssistantMessageCell class]
           forCellReuseIdentifier:kAssistantCellIdentifier];
    [self.tableView registerClass:[KBChatTimeCell class]
           forCellReuseIdentifier:kTimeCellIdentifier];
    
    // 布局
    [self.tableView mas_makeConstraints:^(MASConstraintMaker *make) {
        make.edges.equalTo(self);
    }];
}

#pragma mark - Public Methods

- (void)addUserMessage:(NSString *)text {
    // 记录插入前的消息数量
    NSInteger oldCount = self.messages.count;
    
    KBAiChatMessage *message = [KBAiChatMessage userMessageWithText:text];
    [self insertMessageWithTimestamp:message];
    
    // 计算新增的行数
    NSInteger newCount = self.messages.count;
    NSInteger insertedCount = newCount - oldCount;
    
    // 使用 insert 插入新行
    if (insertedCount > 0) {
        NSMutableArray *indexPaths = [NSMutableArray array];
        for (NSInteger i = oldCount; i < newCount; i++) {
            [indexPaths addObject:[NSIndexPath indexPathForRow:i inSection:0]];
        }
        [self.tableView insertRowsAtIndexPaths:indexPaths
                              withRowAnimation:UITableViewRowAnimationNone];
    }
    
    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.1 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
        [self scrollToBottom];
    });
}

- (void)addAssistantMessage:(NSString *)text
              audioDuration:(NSTimeInterval)duration
                  audioData:(NSData *)audioData {
    // 记录插入前的消息数量
    NSInteger oldCount = self.messages.count;
    
    KBAiChatMessage *message = [KBAiChatMessage assistantMessageWithText:text
                                                        audioDuration:duration
                                                            audioData:audioData];
    [self insertMessageWithTimestamp:message];
    
    // 计算新增的行数
    NSInteger newCount = self.messages.count;
    NSInteger insertedCount = newCount - oldCount;
    
    // 使用 insert 插入新行
    if (insertedCount > 0) {
        NSMutableArray *indexPaths = [NSMutableArray array];
        for (NSInteger i = oldCount; i < newCount; i++) {
            [indexPaths addObject:[NSIndexPath indexPathForRow:i inSection:0]];
        }
        [self.tableView insertRowsAtIndexPaths:indexPaths
                              withRowAnimation:UITableViewRowAnimationNone];
    }
    
    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.1 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
        [self scrollToBottom];
    });
}

- (void)addAssistantMessage:(NSString *)text
                    audioId:(NSString *)audioId {
    NSLog(@"[KBChatTableView] ========== 添加新的 AI 消息 ==========");
    NSLog(@"[KBChatTableView] 文本长度: %lu, audioId: %@", (unsigned long)text.length, audioId);
    NSLog(@"[KBChatTableView] 当前消息数量: %ld", (long)self.messages.count);
    
    // 在添加新消息之前，先标记上一条 AI 消息完成，并停止其打字机效果
    for (NSInteger i = self.messages.count - 1; i >= 0; i--) {
        KBAiChatMessage *msg = self.messages[i];
        if (msg.type == KBAiChatMessageTypeAssistant && !msg.isComplete) {
            NSLog(@"[KBChatTableView] 找到上一条未完成的消息 - 索引: %ld, 文本: %@", (long)i, msg.text);
            msg.isComplete = YES;
            msg.needsTypewriterEffect = NO;
            
            // 停止该 Cell 的打字机效果
            NSIndexPath *oldIndexPath = [NSIndexPath indexPathForRow:i inSection:0];
            KBChatAssistantMessageCell *oldCell = [self.tableView cellForRowAtIndexPath:oldIndexPath];
            if ([oldCell isKindOfClass:[KBChatAssistantMessageCell class]]) {
                NSLog(@"[KBChatTableView] 停止上一条消息的打字机效果");
                [oldCell stopTypewriterEffect];
                // 显示完整文本
                oldCell.messageLabel.text = msg.text;
            }
            break;
        }
    }
    
    // 记录插入前的消息数量
    NSInteger oldCount = self.messages.count;
    
    KBAiChatMessage *message = [KBAiChatMessage assistantMessageWithText:text
                                                              audioId:audioId];
    message.needsTypewriterEffect = YES; // 新消息需要打字机效果
    NSLog(@"[KBChatTableView] 新消息属性 - needsTypewriter: %d, isComplete: %d", 
          message.needsTypewriterEffect, message.isComplete);
    [self insertMessageWithTimestamp:message];
    
    // 计算新增的行数
    NSInteger newCount = self.messages.count;
    NSInteger insertedCount = newCount - oldCount;
    NSLog(@"[KBChatTableView] 插入后消息数量: %ld, 新增行数: %ld", (long)newCount, (long)insertedCount);
    
    // 使用 insert 插入新行
    if (insertedCount > 0) {
        NSMutableArray *indexPaths = [NSMutableArray array];
        for (NSInteger i = oldCount; i < newCount; i++) {
            [indexPaths addObject:[NSIndexPath indexPathForRow:i inSection:0]];
            NSLog(@"[KBChatTableView] 将插入行: %ld", (long)i);
        }
        [self.tableView insertRowsAtIndexPaths:indexPaths
                              withRowAnimation:UITableViewRowAnimationNone];
    }
    
    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.1 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
        [self scrollToBottom];
    });
    
    NSLog(@"[KBChatTableView] ========== 添加完成 ==========");
}

- (void)updateLastAssistantMessage:(NSString *)text {
    // 查找最后一条未完成的 AI 消息
    for (NSInteger i = self.messages.count - 1; i >= 0; i--) {
        KBAiChatMessage *message = self.messages[i];
        if (message.type == KBAiChatMessageTypeAssistant && !message.isComplete) {
            NSLog(@"[KBChatTableView] 更新最后一条 AI 消息 - 索引: %ld, 文本长度: %lu, needsTypewriter: %d", 
                  (long)i, (unsigned long)text.length, message.needsTypewriterEffect);
            message.text = text;
            
            // 直接更新 Cell 的文本，不刷新整个 Cell（避免打断打字机效果）
            NSIndexPath *indexPath = [NSIndexPath indexPathForRow:i inSection:0];
            KBChatAssistantMessageCell *cell = [self.tableView cellForRowAtIndexPath:indexPath];
            if ([cell isKindOfClass:[KBChatAssistantMessageCell class]]) {
                NSLog(@"[KBChatTableView] 找到 Cell，直接配置消息");
                // 直接调用 configureWithMessage，让 Cell 自己决定是否使用打字机效果
                [cell configureWithMessage:message];
            } else {
                NSLog(@"[KBChatTableView] 未找到 Cell 或类型不匹配");
            }
            return;
        }
    }
    
    // 如果没找到，添加新消息
    NSLog(@"[KBChatTableView] 未找到未完成的 AI 消息，添加新消息");
    [self addAssistantMessage:text audioDuration:0 audioData:nil];
}

- (void)markLastAssistantMessageComplete {
    for (NSInteger i = self.messages.count - 1; i >= 0; i--) {
        KBAiChatMessage *message = self.messages[i];
        if (message.type == KBAiChatMessageTypeAssistant) {
            NSLog(@"[KBChatTableView] 标记消息完成 - 索引: %ld, 文本: %@", (long)i, message.text);
            message.isComplete = YES;
            message.needsTypewriterEffect = NO; // 完成后不再需要打字机效果
            
            // 刷新 Cell 以显示完整文本
            NSIndexPath *indexPath = [NSIndexPath indexPathForRow:i inSection:0];
            [self.tableView reloadRowsAtIndexPaths:@[indexPath]
                                  withRowAnimation:UITableViewRowAnimationNone];
            return;
        }
    }
}

- (void)clearMessages {
    [self.messages removeAllObjects];
    [self.tableView reloadData];
}

- (void)scrollToBottom {
    if (self.messages.count == 0) return;
    
    NSIndexPath *lastIndexPath = [NSIndexPath indexPathForRow:self.messages.count - 1
                                                     inSection:0];
    [self.tableView scrollToRowAtIndexPath:lastIndexPath
                          atScrollPosition:UITableViewScrollPositionBottom
                                  animated:YES];
}

#pragma mark - Private Methods

/// 插入消息并自动添加时间戳
- (void)insertMessageWithTimestamp:(KBAiChatMessage *)message {
    // 判断是否需要插入时间戳
    if ([self shouldInsertTimestampForMessage:message]) {
        KBAiChatMessage *timeMessage = [KBAiChatMessage timeMessageWithTimestamp:message.timestamp];
        [self.messages addObject:timeMessage];
    }
    
    [self.messages addObject:message];
}

/// 判断是否需要插入时间戳
- (BOOL)shouldInsertTimestampForMessage:(KBAiChatMessage *)message {
    // 第一条消息总是显示时间
    if (self.messages.count == 0) {
        return YES;
    }
    
    // 查找最后一条非时间戳消息
    KBAiChatMessage *lastMessage = nil;
    for (NSInteger i = self.messages.count - 1; i >= 0; i--) {
        KBAiChatMessage *msg = self.messages[i];
        if (msg.type != KBAiChatMessageTypeTime) {
            lastMessage = msg;
            break;
        }
    }
    
    if (!lastMessage) {
        return YES;
    }
    
    // 计算时间间隔
    NSTimeInterval interval = [message.timestamp timeIntervalSinceDate:lastMessage.timestamp];
    
    // 超过 5 分钟或跨天则显示时间
    if (interval >= kTimestampInterval) {
        return YES;
    }
    
    NSCalendar *calendar = [NSCalendar currentCalendar];
    NSDateComponents *lastComponents = [calendar components:NSCalendarUnitDay | NSCalendarUnitMonth | NSCalendarUnitYear
                                                    fromDate:lastMessage.timestamp];
    NSDateComponents *currentComponents = [calendar components:NSCalendarUnitDay | NSCalendarUnitMonth | NSCalendarUnitYear
                                                       fromDate:message.timestamp];
    
    return ![lastComponents isEqual:currentComponents];
}

/// 刷新并滚动到底部
- (void)reloadAndScroll {
    // 使用 insert 而不是 reload，避免刷新已有的 Cell
    NSInteger lastIndex = self.messages.count - 1;
    if (lastIndex >= 0) {
        NSIndexPath *indexPath = [NSIndexPath indexPathForRow:lastIndex inSection:0];
        [self.tableView insertRowsAtIndexPaths:@[indexPath]
                              withRowAnimation:UITableViewRowAnimationNone];
    }
    
    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.1 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
        [self scrollToBottom];
    });
}

#pragma mark - UITableViewDataSource

- (NSInteger)tableView:(UITableView *)tableView numberOfRowsInSection:(NSInteger)section {
    return self.messages.count;
}

- (UITableViewCell *)tableView:(UITableView *)tableView cellForRowAtIndexPath:(NSIndexPath *)indexPath {
    KBAiChatMessage *message = self.messages[indexPath.row];
    
    NSLog(@"[KBChatTableView] cellForRow: %ld, 消息类型: %ld, needsTypewriter: %d, isComplete: %d", 
          (long)indexPath.row, (long)message.type, message.needsTypewriterEffect, message.isComplete);
    
    switch (message.type) {
        case KBAiChatMessageTypeUser: {
            KBChatUserMessageCell *cell = [tableView dequeueReusableCellWithIdentifier:kUserCellIdentifier
                                                                          forIndexPath:indexPath];
            [cell configureWithMessage:message];
            return cell;
        }
            
        case KBAiChatMessageTypeAssistant: {
            KBChatAssistantMessageCell *cell = [tableView dequeueReusableCellWithIdentifier:kAssistantCellIdentifier
                                                                                forIndexPath:indexPath];
            cell.delegate = self;
            [cell configureWithMessage:message];
            
            // 更新播放状态
            BOOL isPlaying = [indexPath isEqual:self.playingCellIndexPath];
            [cell updateVoicePlayingState:isPlaying];
            
            return cell;
        }
            
        case KBAiChatMessageTypeTime: {
            KBChatTimeCell *cell = [tableView dequeueReusableCellWithIdentifier:kTimeCellIdentifier
                                                                    forIndexPath:indexPath];
            [cell configureWithMessage:message];
            return cell;
        }
    }
}

#pragma mark - KBChatAssistantMessageCellDelegate

- (void)assistantMessageCell:(KBChatAssistantMessageCell *)cell
    didTapVoiceButtonForMessage:(KBAiChatMessage *)message {
    
    NSIndexPath *indexPath = [self.tableView indexPathForCell:cell];
    if (!indexPath) return;
    
    // 如果正在播放同一条消息，则暂停
    if ([indexPath isEqual:self.playingCellIndexPath]) {
        [self stopPlayingAudio];
        return;
    }
    
    // 停止之前的播放
    [self stopPlayingAudio];
    
    // 如果有 audioData，直接播放
    if (message.audioData && message.audioData.length > 0) {
        [self playAudioForMessage:message atIndexPath:indexPath];
        return;
    }
    
    // 如果有 audioId，异步加载音频
    if (message.audioId.length > 0) {
        [self loadAndPlayAudioForMessage:message atIndexPath:indexPath];
        return;
    }
    
    NSLog(@"[KBChatTableView] 没有音频数据或 audioId");
}

#pragma mark - Audio Playback

- (void)loadAndPlayAudioForMessage:(KBAiChatMessage *)message atIndexPath:(NSIndexPath *)indexPath {
    // 显示加载动画
    KBChatAssistantMessageCell *cell = [self.tableView cellForRowAtIndexPath:indexPath];
    if ([cell isKindOfClass:[KBChatAssistantMessageCell class]]) {
        [cell showLoadingAnimation];
    }
    
    // 开始轮询请求（最多5次，每次间隔0.5秒）
    [self pollAudioForMessage:message atIndexPath:indexPath retryCount:0 maxRetries:5];
}

- (void)pollAudioForMessage:(KBAiChatMessage *)message 
                atIndexPath:(NSIndexPath *)indexPath 
                 retryCount:(NSInteger)retryCount 
                 maxRetries:(NSInteger)maxRetries {
    
    __weak typeof(self) weakSelf = self;
    [self.aiVM requestAudioWithAudioId:message.audioId
                            completion:^(NSString *_Nullable audioURL, NSError *_Nullable error) {
        __strong typeof(weakSelf) strongSelf = weakSelf;
        if (!strongSelf) return;
        
        dispatch_async(dispatch_get_main_queue(), ^{
            // 如果成功获取到 audioURL
            if (!error && audioURL.length > 0) {
                // 下载并播放音频
                [strongSelf downloadAndPlayAudioFromURL:audioURL
                                              forMessage:message
                                             atIndexPath:indexPath];
                return;
            }
            
            // 如果还没达到最大重试次数，继续轮询
            if (retryCount < maxRetries - 1) {
                NSLog(@"[KBChatTableView] 音频未就绪，0.5秒后重试 (%ld/%ld)", 
                      (long)(retryCount + 1), (long)maxRetries);
                
                // 0.5秒后重试
                dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.5 * NSEC_PER_SEC)), 
                              dispatch_get_main_queue(), ^{
                    [strongSelf pollAudioForMessage:message 
                                        atIndexPath:indexPath 
                                         retryCount:retryCount + 1 
                                         maxRetries:maxRetries];
                });
            } else {
                // 达到最大重试次数，隐藏加载动画
                KBChatAssistantMessageCell *cell = [strongSelf.tableView cellForRowAtIndexPath:indexPath];
                if ([cell isKindOfClass:[KBChatAssistantMessageCell class]]) {
                    [cell hideLoadingAnimation];
                }
                NSLog(@"[KBChatTableView] 音频加载失败，已重试 %ld 次", (long)maxRetries);
            }
        });
    }];
}

- (void)downloadAndPlayAudioFromURL:(NSString *)urlString
                         forMessage:(KBAiChatMessage *)message
                        atIndexPath:(NSIndexPath *)indexPath {
    NSURL *url = [NSURL URLWithString:urlString];
    if (!url) {
        NSLog(@"[KBChatTableView] 无效的音频 URL: %@", urlString);
        // 隐藏加载动画
        KBChatAssistantMessageCell *cell = [self.tableView cellForRowAtIndexPath:indexPath];
        if ([cell isKindOfClass:[KBChatAssistantMessageCell class]]) {
            [cell hideLoadingAnimation];
        }
        return;
    }
    
    NSURLSessionConfiguration *config = [NSURLSessionConfiguration defaultSessionConfiguration];
    NSURLSession *session = [NSURLSession sessionWithConfiguration:config];
    
    NSURLSessionDataTask *task = [session dataTaskWithURL:url
                                        completionHandler:^(NSData *_Nullable data,
                                                            NSURLResponse *_Nullable response,
                                                            NSError *_Nullable error) {
        if (error || !data || data.length == 0) {
            NSLog(@"[KBChatTableView] 下载音频失败: %@", error.localizedDescription ?: @"");
            dispatch_async(dispatch_get_main_queue(), ^{
                // 隐藏加载动画
                KBChatAssistantMessageCell *cell = [self.tableView cellForRowAtIndexPath:indexPath];
                if ([cell isKindOfClass:[KBChatAssistantMessageCell class]]) {
                    [cell hideLoadingAnimation];
                }
            });
            return;
        }
        
        dispatch_async(dispatch_get_main_queue(), ^{
            // 隐藏加载动画
            KBChatAssistantMessageCell *cell = [self.tableView cellForRowAtIndexPath:indexPath];
            if ([cell isKindOfClass:[KBChatAssistantMessageCell class]]) {
                [cell hideLoadingAnimation];
            }
            
            // 缓存音频数据到消息对象
            message.audioData = data;
            
            // 计算音频时长
            NSError *playerError = nil;
            AVAudioPlayer *player = [[AVAudioPlayer alloc] initWithData:data error:&playerError];
            if (!playerError && player) {
                message.audioDuration = player.duration;
            }
            
            // 不刷新 Cell，避免触发打字机效果
            // 直接播放音频
            [self playAudioForMessage:message atIndexPath:indexPath];
        });
    }];
    
    [task resume];
}

- (void)playAudioForMessage:(KBAiChatMessage *)message atIndexPath:(NSIndexPath *)indexPath {
    if (!message.audioData || message.audioData.length == 0) {
        NSLog(@"[KBChatTableView] 没有音频数据");
        return;
    }
    
    // 配置音频会话为播放模式
    NSError *sessionError = nil;
    AVAudioSession *audioSession = [AVAudioSession sharedInstance];
    [audioSession setCategory:AVAudioSessionCategoryPlayback error:&sessionError];
    [audioSession setActive:YES error:&sessionError];
    
    if (sessionError) {
        NSLog(@"[KBChatTableView] 音频会话配置失败: %@", sessionError.localizedDescription);
    }
    
    NSError *error = nil;
    self.audioPlayer = [[AVAudioPlayer alloc] initWithData:message.audioData error:&error];
    
    if (error || !self.audioPlayer) {
        NSLog(@"[KBChatTableView] 音频播放器初始化失败: %@", error.localizedDescription);
        return;
    }
    
    self.audioPlayer.delegate = self;
    self.audioPlayer.volume = 1.0; // 设置音量为最大
    [self.audioPlayer prepareToPlay];
    [self.audioPlayer play];
    
    self.playingCellIndexPath = indexPath;
    
    // 更新 Cell 状态
    KBChatAssistantMessageCell *cell = [self.tableView cellForRowAtIndexPath:indexPath];
    if ([cell isKindOfClass:[KBChatAssistantMessageCell class]]) {
        [cell updateVoicePlayingState:YES];
    }
}

- (void)stopPlayingAudio {
    if (self.audioPlayer && self.audioPlayer.isPlaying) {
        [self.audioPlayer stop];
    }
    
    if (self.playingCellIndexPath) {
        KBChatAssistantMessageCell *cell = [self.tableView cellForRowAtIndexPath:self.playingCellIndexPath];
        if ([cell isKindOfClass:[KBChatAssistantMessageCell class]]) {
            [cell updateVoicePlayingState:NO];
        }
        self.playingCellIndexPath = nil;
    }
}

#pragma mark - AVAudioPlayerDelegate

- (void)audioPlayerDidFinishPlaying:(AVAudioPlayer *)player successfully:(BOOL)flag {
    [self stopPlayingAudio];
}

@end