keyboard/CustomKeyboard/Network/KBStreamFetcher.m

//
//  KBStreamFetcher.m
//

#import "KBStreamFetcher.h"

@interface KBStreamFetcher () <NSURLSessionDataDelegate>
@property (nonatomic, strong) NSURLSession *session;
@property (nonatomic, strong) NSURLSessionDataTask *task;
@property (nonatomic, strong) NSMutableData *buffer;           // 网络原始字节累加
@property (nonatomic, assign) NSStringEncoding textEncoding;   // 推断得到的文本编码（默认 UTF-8）
@property (nonatomic, assign) BOOL isSSE;                      // 是否为 SSE 响应
@property (nonatomic, strong) NSMutableString *sseTextBuffer;  // SSE 文本缓冲（已解码）
@property (nonatomic, assign) NSInteger decodedPrefixBytes;    // 已解码并写入 sseTextBuffer 的字节数（SSE）
@property (nonatomic, assign) NSInteger deliveredCharCount;    // 已回传的字符数（非 SSE，用于做增量）
@property (nonatomic, assign) BOOL hasEmitted;                 // 是否已经输出过正文（用于“首段删 1 个 \t”）
@property (nonatomic, assign) BOOL lastChunkEndedWithTab;      // 上一个已输出分片是否以 "\t" 结尾（用于跨分片去除“\t 后空格”）
@property (nonatomic, strong) NSMutableArray<NSString *> *pendingQueue; // 待回调的分片（节流输出）
@property (nonatomic, strong) NSTimer *flushTimer;             // 定时从队列取出一条回调
@property (nonatomic, strong, nullable) NSError *finishError;  // 结束时的错误（需要等队列清空再回调）

// Metrics
@property (nonatomic, assign) CFAbsoluteTime tStart;           // start() 被调用的时刻
@property (nonatomic, assign) CFAbsoluteTime tFirstByte;       // 第一次拿到可解码内容
@property (nonatomic, assign) CFAbsoluteTime tFinish;          // 完成/失败时刻
@property (nonatomic, assign) NSInteger emittedChunkCount;     // 已输出分片数量
@end

// 计算数据中以 UTF-8 编码可完整解码的“前缀字节长度”，避免切断多字节字符
static NSUInteger kb_validUTF8PrefixLen(NSData *data) {
    const unsigned char *bytes = (const unsigned char *)data.bytes;
    NSUInteger n = data.length;
    if (n == 0) return 0;
    NSInteger i = (NSInteger)n - 1;
    while (i >= 0 && (bytes[i] & 0xC0) == 0x80) { i--; } // 10xxxxxx 续字节
    if (i < 0) return 0; // 全是续字节，等下次
    unsigned char b = bytes[i];
    NSUInteger expected = 1;
    if ((b & 0x80) == 0x00) expected = 1;        // 0xxxxxxx
    else if ((b & 0xE0) == 0xC0) expected = 2;   // 110xxxxx
    else if ((b & 0xF0) == 0xE0) expected = 3;   // 1110xxxx
    else if ((b & 0xF8) == 0xF0) expected = 4;   // 11110xxx
    else return (NSUInteger)i;                   // 非法起始，截到 i 之前
    NSUInteger remain = n - (NSUInteger)i;
    return (remain >= expected) ? n : (NSUInteger)i;
}

@implementation KBStreamFetcher

+ (instancetype)fetcherWithURL:(NSURL *)url {
    KBStreamFetcher *f = [[self alloc] init];
    f.url = url;
    return f;
}

- (instancetype)init {
    if (self = [super init]) {
        _acceptEventStream = NO;
        _disableCompression = YES;
        _treatSlashTAsTab = YES;
        _trimLeadingTabOnce = YES;
        _requestTimeout = 30.0;
        _textEncoding = NSUTF8StringEncoding;
        _buffer = [NSMutableData data];
        _sseTextBuffer = [NSMutableString string];
        _pendingQueue = [NSMutableArray array];
        _flushInterval = 0.1;
        _splitLargeDeltasOnWhitespace = YES;
        _loggingEnabled = YES;
    }
    return self;
}

- (void)start {
    if (!self.url) return;
    [self cancel];

    NSURLSessionConfiguration *cfg = [NSURLSessionConfiguration defaultSessionConfiguration];
    cfg.requestCachePolicy = NSURLRequestReloadIgnoringLocalCacheData;
    cfg.timeoutIntervalForRequest = self.requestTimeout;
    cfg.timeoutIntervalForResource = MAX(self.requestTimeout, 60.0);

    self.session = [NSURLSession sessionWithConfiguration:cfg delegate:self delegateQueue:[NSOperationQueue mainQueue]];

    NSMutableURLRequest *req = [NSMutableURLRequest requestWithURL:self.url];
    req.HTTPMethod = @"GET";
    if (self.disableCompression) { [req setValue:@"identity" forHTTPHeaderField:@"Accept-Encoding"]; }
    if (self.acceptEventStream) { [req setValue:@"text/event-stream" forHTTPHeaderField:@"Accept"]; }
    [req setValue:@"no-cache" forHTTPHeaderField:@"Cache-Control"];
    [req setValue:@"keep-alive" forHTTPHeaderField:@"Connection"];
    [self.extraHeaders enumerateKeysAndObjectsUsingBlock:^(NSString *k, NSString *v, BOOL *stop){ [req setValue:v forHTTPHeaderField:k]; }];

    // 状态复位
    [self.buffer setLength:0];
    [self.sseTextBuffer setString:@""];
    self.isSSE = NO;
    self.textEncoding = NSUTF8StringEncoding;
    self.decodedPrefixBytes = 0;
    self.deliveredCharCount = 0;
    self.hasEmitted = NO;
    self.lastChunkEndedWithTab = NO;
    [self.pendingQueue removeAllObjects];
    [self.flushTimer invalidate]; self.flushTimer = nil;
    self.finishError = nil;

    self.tStart = CFAbsoluteTimeGetCurrent();
    self.tFirstByte = 0;
    self.tFinish = 0;
    self.emittedChunkCount = 0;
    if (self.loggingEnabled) {
        NSLog(@"[KBStream] start url=%@ acceptSSE=%@ disableCompression=%@ flush=%.0fms splitWords=%@",
              self.url.absoluteString,
              self.acceptEventStream?@"YES":@"NO",
              self.disableCompression?@"YES":@"NO",
              self.flushInterval*1000.0,
              self.splitLargeDeltasOnWhitespace?@"YES":@"NO");
    }
    self.task = [self.session dataTaskWithRequest:req];
    [self.task resume];
}

- (void)cancel {
    [self.task cancel];
    self.task = nil;
    [self.session invalidateAndCancel];
    self.session = nil;
    [self.buffer setLength:0];
    [self.sseTextBuffer setString:@""];
    self.decodedPrefixBytes = 0;
    self.deliveredCharCount = 0;
    self.hasEmitted = NO;
    self.lastChunkEndedWithTab = NO;
    [self.pendingQueue removeAllObjects];
    [self.flushTimer invalidate]; self.flushTimer = nil;
    self.finishError = nil;
}

#pragma mark - NSURLSessionDataDelegate

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveResponse:(NSURLResponse *)response completionHandler:(void (^)(NSURLSessionResponseDisposition))completionHandler {
    self.isSSE = NO;
    self.textEncoding = NSUTF8StringEncoding;
    if ([response isKindOfClass:[NSHTTPURLResponse class]]) {
        NSHTTPURLResponse *r = (NSHTTPURLResponse *)response;
        NSString *ct = r.allHeaderFields[@"Content-Type"] ?: r.allHeaderFields[@"content-type"];
        if ([ct isKindOfClass:[NSString class]]) {
            NSString *lower = [ct lowercaseString];
            if ([lower containsString:@"text/event-stream"]) self.isSSE = YES;
            NSRange pos = [lower rangeOfString:@"charset="];
            if (pos.location != NSNotFound) {
                NSString *charset = [[lower substringFromIndex:pos.location + pos.length] componentsSeparatedByString:@";"][0];
                if ([charset containsString:@"utf-8"] || [charset containsString:@"utf8"]) {
                    self.textEncoding = NSUTF8StringEncoding;
                } else if ([charset containsString:@"iso-8859-1"] || [charset containsString:@"latin1"]) {
                    self.textEncoding = NSISOLatin1StringEncoding;
                }
            }
        }
    }
    [self.sseTextBuffer setString:@""];
    self.decodedPrefixBytes = 0;
    if (completionHandler) completionHandler(NSURLSessionResponseAllow);
}

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveData:(NSData *)data {
    if (data.length == 0) return;
    [self.buffer appendData:data];

    NSUInteger validLen = (self.textEncoding == NSUTF8StringEncoding)
        ? kb_validUTF8PrefixLen(self.buffer)
        : self.buffer.length;
    if (validLen > 0 && self.tFirstByte == 0) {
        self.tFirstByte = CFAbsoluteTimeGetCurrent();
        if (self.loggingEnabled) {
            NSLog(@"[KBStream] first-bytes after %.0fms (encoding=%@, SSE=%@)",
                  (self.tFirstByte - self.tStart)*1000.0,
                  (self.textEncoding==NSUTF8StringEncoding?@"UTF-8":@"Other"),
                  self.isSSE?@"YES":@"NO");
        }
    }
    if (validLen == 0) return; // 末尾可能卡着半个字符

    if (self.isSSE) {
        if ((NSUInteger)self.decodedPrefixBytes < validLen) {
            NSRange rng = NSMakeRange((NSUInteger)self.decodedPrefixBytes, validLen - (NSUInteger)self.decodedPrefixBytes);
            NSString *piece = [[NSString alloc] initWithBytes:(const char *)self.buffer.bytes + rng.location
                                                       length:rng.length
                                                     encoding:self.textEncoding];
            if (piece.length > 0) {
                [self.sseTextBuffer appendString:piece];
                self.decodedPrefixBytes = (NSInteger)validLen;
            }
        }
        // 统一换行并按 SSE 事件 \n\n 切开
        if (self.sseTextBuffer.length > 0) {
            NSString *normalized = [self.sseTextBuffer stringByReplacingOccurrencesOfString:@"\r\n" withString:@"\n"];
            [self.sseTextBuffer setString:normalized];
            while (1) {
                NSRange sep = [self.sseTextBuffer rangeOfString:@"\n\n"]; // 完整事件
                if (sep.location == NSNotFound) break;
                NSString *event = [self.sseTextBuffer substringToIndex:sep.location];
                [self.sseTextBuffer deleteCharactersInRange:NSMakeRange(0, sep.location + sep.length)];

                // 合并 data: 行为正文
                NSArray<NSString *> *lines = [event componentsSeparatedByString:@"\n"];
                NSMutableString *payload = [NSMutableString string];
                for (NSString *ln in lines) {
                    if ([ln hasPrefix:@"data:"]) {
                        NSString *v = [ln substringFromIndex:5];
                        if (v.length > 0 && [v hasPrefix:@" "]) v = [v substringFromIndex:1];
                        [payload appendString:v ?: @""];
                    }
                }
                if (payload.length > 0) { [self enqueueChunk:payload]; }
            }
        }
        return;
    }

    // 非 SSE：直接对“可解码前缀”做增量输出
    NSString *prefix = [[NSString alloc] initWithBytes:self.buffer.bytes length:validLen encoding:self.textEncoding];
    if (!prefix) return;
    if (self.deliveredCharCount < (NSInteger)prefix.length) {
        NSString *delta = [prefix substringFromIndex:self.deliveredCharCount];
        self.deliveredCharCount = prefix.length;
        if (self.splitLargeDeltasOnWhitespace && delta.length > 16) {
            // 按空格切词逐条回调（保留空格，使观感更自然）
            NSArray<NSString *> *parts = [delta componentsSeparatedByString:@" "];
            for (NSUInteger i = 0; i < parts.count; i++) {
                NSString *w = parts[i];
                if (w.length == 0) { [self enqueueChunk:@" "]; continue; }
                if (i + 1 < parts.count) {
                    [self enqueueChunk:[w stringByAppendingString:@" "]];
                } else {
                    [self enqueueChunk:w];
                }
            }
        } else {
            [self enqueueChunk:delta];
        }
    }
}

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didCompleteWithError:(NSError *)error {
    if (!error && self.isSSE && self.sseTextBuffer.length > 0) {
        // 处理最后一条未以 \n\n 结束的事件
        NSString *normalized = [self.sseTextBuffer stringByReplacingOccurrencesOfString:@"\r\n" withString:@"\n"];
        NSArray<NSString *> *lines = [normalized componentsSeparatedByString:@"\n"];
        NSMutableString *payload = [NSMutableString string];
        for (NSString *ln in lines) {
            if ([ln hasPrefix:@"data:"]) {
                NSString *v = [ln substringFromIndex:5];
                if (v.length > 0 && [v hasPrefix:@" "]) v = [v substringFromIndex:1];
                [payload appendString:v ?: @""];
            }
        }
        if (payload.length > 0) {
            NSString *delta = nil;
            if ((NSInteger)payload.length >= self.deliveredCharCount) {
                delta = [payload substringFromIndex:self.deliveredCharCount];
            } else {
                delta = payload;
            }
            self.deliveredCharCount = payload.length;
            if (delta.length > 0) { [self emitChunk:delta]; }
        }
    }

    self.tFinish = CFAbsoluteTimeGetCurrent();
    if (self.loggingEnabled) {
        double t0 = (self.tFirstByte>0? (self.tFirstByte - self.tStart)*1000.0 : -1);
        double t1 = (self.tFirstByte>0? (self.tFinish - self.tFirstByte)*1000.0 : -1);
        double tt = (self.tFinish - self.tStart)*1000.0;
        NSLog(@"[KBStream] finish chunks=%ld firstByte=%.0fms after start, tail=%.0fms, total=%.0fms error=%@",
              (long)self.emittedChunkCount, t0, t1, tt, error);
    }
    // 若队列还有待输出内容，等队列清空再回调 finish
    if (self.pendingQueue.count > 0) {
        self.finishError = error;
        [self startFlushTimerIfNeeded];
    } else {
        if (self.onFinish) dispatch_async(dispatch_get_main_queue(), ^{ self.onFinish(error); });
        [self cancel];
    }
}

#pragma mark - Helpers

- (void)emitChunk:(NSString *)rawText {
    if (rawText.length == 0) return;
    NSString *text = rawText;
    // 0) 规范化换行与段起始：去掉位于片段开头的 \r/\n；将 "\n\t"、"\r\n\t"、"\r\t" 归一为 "\t"
    text = [text stringByReplacingOccurrencesOfString:@"\r\n\t" withString:@"\t"];
    text = [text stringByReplacingOccurrencesOfString:@"\n\t" withString:@"\t"];
    text = [text stringByReplacingOccurrencesOfString:@"\r\t" withString:@"\t"];
    while (text.length > 0) {
        unichar c0 = [text characterAtIndex:0];
        if (c0 == '\n' || c0 == '\r') { text = [text substringFromIndex:1]; continue; }
        break;
    }
    // 1) 统一处理 “/t” -> “\t”
    if (self.treatSlashTAsTab) {
        text = [text stringByReplacingOccurrencesOfString:@"/t" withString:@"\t"];
    }
    // 2) 仅在整体首段：去掉一个起始的 "\t"，以及其后紧邻的一个空格（若存在）
    if (!self.hasEmitted && self.trimLeadingTabOnce) {
        if (text.length > 0 && [text characterAtIndex:0] == '\t') {
            NSUInteger start = 1;
            if (start < text.length && [text characterAtIndex:start] == ' ') start++;
            text = [text substringFromIndex:start];
        }
    }
    // 3) 从第二段开始：去掉每个段首的一个空格（即 “\t ” -> “\t”），跨分片也处理
    if (text.length > 0) {
        // 跨分片：若上个分片以 \t 结尾，本分片起始的一个或多个空格去掉一个
        if (self.lastChunkEndedWithTab) {
            NSUInteger j = 0;
            while (j < text.length && [text characterAtIndex:j] == ' ') { j++; }
            if (j > 0) {
                text = [text substringFromIndex:1]; // 仅去一个空格
            }
        }
        // 同一分片内：将 “\t ” 规范化为 “\t”（仅去一个空格）
        text = [text stringByReplacingOccurrencesOfString:@"\t " withString:@"\t"];
    }
    if (text.length == 0) { self.lastChunkEndedWithTab = NO; return; }
    self.emittedChunkCount += 1;
    if (self.loggingEnabled) {
        NSLog(@"[KBStream] chunk#%ld len=%lu text=\"%@\"",
              (long)self.emittedChunkCount, (unsigned long)text.length, KBPrintableSnippet(text, 160));
    }
    if (self.onChunk) dispatch_async(dispatch_get_main_queue(), ^{ self.onChunk(text); });
    self.hasEmitted = YES;
    // 记录末尾是否为分段分隔符 \t（用于跨分片处理）
    unichar lastc = [text characterAtIndex:text.length - 1];
    self.lastChunkEndedWithTab = (lastc == '\t');
}

#pragma mark - Queue/Flush

- (void)enqueueChunk:(NSString *)s {
    if (s.length == 0) return;
    [self.pendingQueue addObject:s];
    [self startFlushTimerIfNeeded];
}

- (void)startFlushTimerIfNeeded {
    if (self.flushTimer) return;
    __weak typeof(self) weakSelf = self;
    self.flushTimer = [NSTimer scheduledTimerWithTimeInterval:MAX(0.01, self.flushInterval)
                                                     repeats:YES
                                                       block:^(NSTimer * _Nonnull t) {
        __strong typeof(weakSelf) self = weakSelf; if (!self) { [t invalidate]; return; }
        if (self.pendingQueue.count == 0) {
            [t invalidate]; self.flushTimer = nil;
            if (self.finishError || self.finishError == nil) {
                NSError *err = self.finishError; self.finishError = nil;
                if (self.onFinish) dispatch_async(dispatch_get_main_queue(), ^{ self.onFinish(err); });
                [self cancel];
            }
            return;
        }
        NSString *first = self.pendingQueue.firstObject;
        [self.pendingQueue removeObjectAtIndex:0];
        [self emitChunk:first];
    }];
}

#pragma mark - Logging helpers

static NSString *KBPrintableSnippet(NSString *s, NSUInteger maxLen) {
    if (!s) return @"";
    NSString *x = [s stringByReplacingOccurrencesOfString:@"\n" withString:@"\\n"];
    if (x.length > maxLen) {
        x = [[x substringToIndex:maxLen] stringByAppendingString:@"…"];
    }
    return x;
}

@end