keyboard/keyBoard/Class/AiTalk/VM/ASRStreamClient.m

//
//  ASRStreamClient.m
//  keyBoard
//
//  Created by Mac on 2026/1/15.
//

#import "ASRStreamClient.h"
#import "AudioCaptureManager.h"

@interface ASRStreamClient () <NSURLSessionWebSocketDelegate>

@property(nonatomic, strong) NSURLSession *urlSession;
@property(nonatomic, strong) NSURLSessionWebSocketTask *webSocketTask;
@property(nonatomic, copy) NSString *currentSessionId;
@property(nonatomic, strong) dispatch_queue_t networkQueue;
@property(nonatomic, assign) BOOL connected;

@end

@implementation ASRStreamClient

- (instancetype)init {
  self = [super init];
  if (self) {
    _networkQueue = dispatch_queue_create("com.keyboard.aitalk.asr.network",
                                          DISPATCH_QUEUE_SERIAL);
    // TODO: 替换为实际的 ASR 服务器地址
    _serverURL = @"wss://your-asr-server.com/ws/asr";
  }
  return self;
}

- (void)dealloc {
  [self cancelInternal];
}

#pragma mark - Public Methods

- (void)startWithSessionId:(NSString *)sessionId {
  dispatch_async(self.networkQueue, ^{
    [self cancelInternal];

    self.currentSessionId = sessionId;

    // 创建 WebSocket 连接
    NSURL *url = [NSURL URLWithString:self.serverURL];
    NSURLSessionConfiguration *config =
        [NSURLSessionConfiguration defaultSessionConfiguration];
    config.timeoutIntervalForRequest = 30;
    config.timeoutIntervalForResource = 300;

    self.urlSession = [NSURLSession sessionWithConfiguration:config
                                                    delegate:self
                                               delegateQueue:nil];

    self.webSocketTask = [self.urlSession webSocketTaskWithURL:url];
    [self.webSocketTask resume];

    // 发送 start 消息
    NSDictionary *startMessage = @{
      @"type" : @"start",
      @"sessionId" : sessionId,
      @"format" : @"pcm_s16le",
      @"sampleRate" : @(kAudioSampleRate),
      @"channels" : @(kAudioChannels)
    };

    NSError *jsonError = nil;
    NSData *jsonData = [NSJSONSerialization dataWithJSONObject:startMessage
                                                       options:0
                                                         error:&jsonError];
    if (jsonError) {
      [self reportError:jsonError];
      return;
    }

    NSString *jsonString = [[NSString alloc] initWithData:jsonData
                                                 encoding:NSUTF8StringEncoding];
    NSURLSessionWebSocketMessage *message =
        [[NSURLSessionWebSocketMessage alloc] initWithString:jsonString];

    [self.webSocketTask
              sendMessage:message
        completionHandler:^(NSError *_Nullable error) {
          if (error) {
            [self reportError:error];
          } else {
            self.connected = YES;
            [self receiveMessage];
            NSLog(@"[ASRStreamClient] Started session: %@", sessionId);
          }
        }];
  });
}

- (void)sendAudioPCMFrame:(NSData *)pcmFrame {
  if (!self.connected || !self.webSocketTask) {
    return;
  }

  dispatch_async(self.networkQueue, ^{
    NSURLSessionWebSocketMessage *message =
        [[NSURLSessionWebSocketMessage alloc] initWithData:pcmFrame];
    [self.webSocketTask sendMessage:message
                  completionHandler:^(NSError *_Nullable error) {
                    if (error) {
                      NSLog(@"[ASRStreamClient] Failed to send audio frame: %@",
                            error.localizedDescription);
                    }
                  }];
  });
}

- (void)finalize {
  if (!self.connected || !self.webSocketTask) {
    return;
  }

  dispatch_async(self.networkQueue, ^{
    NSDictionary *finalizeMessage =
        @{@"type" : @"finalize", @"sessionId" : self.currentSessionId ?: @""};

    NSError *jsonError = nil;
    NSData *jsonData = [NSJSONSerialization dataWithJSONObject:finalizeMessage
                                                       options:0
                                                         error:&jsonError];
    if (jsonError) {
      [self reportError:jsonError];
      return;
    }

    NSString *jsonString = [[NSString alloc] initWithData:jsonData
                                                 encoding:NSUTF8StringEncoding];
    NSURLSessionWebSocketMessage *message =
        [[NSURLSessionWebSocketMessage alloc] initWithString:jsonString];

    [self.webSocketTask sendMessage:message
                  completionHandler:^(NSError *_Nullable error) {
                    if (error) {
                      [self reportError:error];
                    } else {
                      NSLog(@"[ASRStreamClient] Sent finalize for session: %@",
                            self.currentSessionId);
                    }
                  }];
  });
}

- (void)cancel {
  dispatch_async(self.networkQueue, ^{
    [self cancelInternal];
  });
}

#pragma mark - Private Methods

- (void)cancelInternal {
  self.connected = NO;

  if (self.webSocketTask) {
    [self.webSocketTask cancel];
    self.webSocketTask = nil;
  }

  if (self.urlSession) {
    [self.urlSession invalidateAndCancel];
    self.urlSession = nil;
  }

  self.currentSessionId = nil;
}

- (void)receiveMessage {
  if (!self.webSocketTask) {
    return;
  }

  __weak typeof(self) weakSelf = self;
  [self.webSocketTask receiveMessageWithCompletionHandler:^(
                          NSURLSessionWebSocketMessage *_Nullable message,
                          NSError *_Nullable error) {
    __strong typeof(weakSelf) strongSelf = weakSelf;
    if (!strongSelf)
      return;

    if (error) {
      // 检查是否是正常关闭
      if (error.code != 57 && error.code != NSURLErrorCancelled) {
        [strongSelf reportError:error];
      }
      return;
    }

    if (message.type == NSURLSessionWebSocketMessageTypeString) {
      [strongSelf handleTextMessage:message.string];
    }

    // 继续接收下一条消息
    [strongSelf receiveMessage];
  }];
}

- (void)handleTextMessage:(NSString *)text {
  NSData *data = [text dataUsingEncoding:NSUTF8StringEncoding];
  NSError *jsonError = nil;
  NSDictionary *json = [NSJSONSerialization JSONObjectWithData:data
                                                       options:0
                                                         error:&jsonError];

  if (jsonError) {
    NSLog(@"[ASRStreamClient] Failed to parse message: %@", text);
    return;
  }

  NSString *type = json[@"type"];

  if ([type isEqualToString:@"partial"]) {
    NSString *partialText = json[@"text"] ?: @"";
    dispatch_async(dispatch_get_main_queue(), ^{
      if ([self.delegate
              respondsToSelector:@selector(asrClientDidReceivePartialText:)]) {
        [self.delegate asrClientDidReceivePartialText:partialText];
      }
    });
  } else if ([type isEqualToString:@"final"]) {
    NSString *finalText = json[@"text"] ?: @"";
    dispatch_async(dispatch_get_main_queue(), ^{
      if ([self.delegate
              respondsToSelector:@selector(asrClientDidReceiveFinalText:)]) {
        [self.delegate asrClientDidReceiveFinalText:finalText];
      }
    });
    // 收到最终结果后关闭连接
    [self cancelInternal];
  } else if ([type isEqualToString:@"error"]) {
    NSInteger code = [json[@"code"] integerValue];
    NSString *message = json[@"message"] ?: @"Unknown error";
    NSError *error =
        [NSError errorWithDomain:@"ASRStreamClient"
                            code:code
                        userInfo:@{NSLocalizedDescriptionKey : message}];
    [self reportError:error];
  }
}

- (void)reportError:(NSError *)error {
  dispatch_async(dispatch_get_main_queue(), ^{
    if ([self.delegate respondsToSelector:@selector(asrClientDidFail:)]) {
      [self.delegate asrClientDidFail:error];
    }
  });
}

#pragma mark - NSURLSessionWebSocketDelegate

- (void)URLSession:(NSURLSession *)session
          webSocketTask:(NSURLSessionWebSocketTask *)webSocketTask
    didOpenWithProtocol:(NSString *)protocol {
  NSLog(@"[ASRStreamClient] WebSocket connected with protocol: %@", protocol);
}

- (void)URLSession:(NSURLSession *)session
       webSocketTask:(NSURLSessionWebSocketTask *)webSocketTask
    didCloseWithCode:(NSURLSessionWebSocketCloseCode)closeCode
              reason:(NSData *)reason {
  NSLog(@"[ASRStreamClient] WebSocket closed with code: %ld", (long)closeCode);
  self.connected = NO;
}

@end