//
//  DeepgramWebSocketClient.m
//  keyBoard
//
//  Created by Mac on 2026/1/21.
//

#import "DeepgramWebSocketClient.h"

static NSString *const kDeepgramWebSocketClientErrorDomain =
    @"DeepgramWebSocketClient";

@interface DeepgramWebSocketClient () <NSURLSessionWebSocketDelegate>

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

@end

@implementation DeepgramWebSocketClient

- (instancetype)init {
  self = [super init];
  if (self) {
    _networkQueue = dispatch_queue_create("com.keyboard.aitalk.deepgram.ws",
                                          DISPATCH_QUEUE_SERIAL);
    _serverURL = @"wss://api.deepgram.com/v1/listen";
    _encoding = @"linear16";
    _sampleRate = 16000.0;
    _channels = 1;
    _punctuate = YES;
    _smartFormat = YES;
    _interimResults = YES;
    _audioSendingEnabled = NO;
  }
  return self;
}

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

#pragma mark - Public Methods

- (void)connect {
  dispatch_async(self.networkQueue, ^{
    [self disconnectInternal];

    if (self.apiKey.length == 0) {
      [self reportErrorWithMessage:@"Deepgram API key is required"];
      return;
    }

    NSURL *url = [self buildURL];
    if (!url) {
      [self reportErrorWithMessage:@"Invalid Deepgram URL"];
      return;
    }

    NSLog(@"[DeepgramWebSocketClient] Connecting: %@", url.absoluteString);

    NSURLSessionConfiguration *config =
        [NSURLSessionConfiguration defaultSessionConfiguration];
    config.timeoutIntervalForRequest = 30;
    config.timeoutIntervalForResource = 300;

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

    NSMutableURLRequest *request = [NSMutableURLRequest requestWithURL:url];
    [request setValue:[NSString stringWithFormat:@"Token %@", self.apiKey]
        forHTTPHeaderField:@"Authorization"];

    self.webSocketTask = [self.urlSession webSocketTaskWithRequest:request];
    [self.webSocketTask resume];
    [self receiveMessage];
  });
}

- (void)disconnect {
  dispatch_async(self.networkQueue, ^{
    BOOL shouldNotify = self.webSocketTask != nil;
    if (shouldNotify) {
      NSLog(@"[DeepgramWebSocketClient] Disconnect requested");
    }
    [self disconnectInternal];
    if (shouldNotify) {
      [self notifyDisconnect:nil];
    }
  });
}

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

  dispatch_async(self.networkQueue, ^{
    if (!self.audioSendingEnabled) {
      return;
    }
    if (!self.connected || !self.webSocketTask) {
      return;
    }

    NSURLSessionWebSocketMessage *message =
        [[NSURLSessionWebSocketMessage alloc] initWithData:pcmFrame];
    [self.webSocketTask
              sendMessage:message
        completionHandler:^(NSError *_Nullable error) {
          if (error) {
            [self reportError:error];
          } else {
            NSLog(@"[DeepgramWebSocketClient] Sent audio frame: %lu bytes",
                  (unsigned long)pcmFrame.length);
          }
        }];
  });
}

- (void)finish {
  NSLog(@"[DeepgramWebSocketClient] Sending CloseStream");
  [self sendJSON:@{@"type" : @"CloseStream"}];
}

- (void)sendKeepAlive {
  if (!self.connected || !self.webSocketTask) {
    return;
  }
  [self sendJSON:@{@"type" : @"KeepAlive"}];
}

- (void)enableAudioSending {
  dispatch_async(self.networkQueue, ^{
    self.audioSendingEnabled = YES;
  });
}

- (void)disableAudioSending {
  dispatch_async(self.networkQueue, ^{
    self.audioSendingEnabled = NO;
  });
}

#pragma mark - Private Methods

- (NSURL *)buildURL {
  if (self.serverURL.length == 0) {
    return nil;
  }

  NSURLComponents *components =
      [NSURLComponents componentsWithString:self.serverURL];
  if (!components) {
    return nil;
  }

  NSMutableArray<NSURLQueryItem *> *items =
      components.queryItems.mutableCopy ?: [NSMutableArray array];

  [self upsertQueryItemWithName:@"model" value:self.model items:items];
  [self upsertQueryItemWithName:@"language" value:self.language items:items];

  [self
      upsertQueryItemWithName:@"punctuate"
                        value:(self.punctuate ? @"true" : @"false")items:items];
  [self upsertQueryItemWithName:@"smart_format"
                          value:(self.smartFormat ? @"true" : @"false")items
                               :items];
  [self upsertQueryItemWithName:@"interim_results"
                          value:(self.interimResults ? @"true" : @"false")items
                               :items];

  [self upsertQueryItemWithName:@"encoding" value:self.encoding items:items];
  [self upsertQueryItemWithName:@"sample_rate"
                          value:[NSString
                                    stringWithFormat:@"%.0f", self.sampleRate]
                          items:items];
  [self upsertQueryItemWithName:@"channels"
                          value:[NSString stringWithFormat:@"%d", self.channels]
                          items:items];

  components.queryItems = items;
  return components.URL;
}

- (void)upsertQueryItemWithName:(NSString *)name
                          value:(NSString *)value
                          items:(NSMutableArray<NSURLQueryItem *> *)items {
  if (name.length == 0 || value.length == 0) {
    return;
  }

  for (NSUInteger i = 0; i < items.count; i++) {
    NSURLQueryItem *item = items[i];
    if ([item.name isEqualToString:name]) {
      items[i] = [NSURLQueryItem queryItemWithName:name value:value];
      return;
    }
  }

  [items addObject:[NSURLQueryItem queryItemWithName:name value:value]];
}

- (void)sendJSON:(NSDictionary *)dict {
  if (!self.webSocketTask) {
    return;
  }

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

  NSString *jsonString = [[NSString alloc] initWithData:jsonData
                                               encoding:NSUTF8StringEncoding];
  if (!jsonString) {
    [self reportErrorWithMessage:@"Failed to encode JSON message"];
    return;
  }

  dispatch_async(self.networkQueue, ^{
    NSURLSessionWebSocketMessage *message =
        [[NSURLSessionWebSocketMessage alloc] initWithString:jsonString];
    [self.webSocketTask sendMessage:message
                  completionHandler:^(NSError *_Nullable error) {
                    if (error) {
                      [self reportError:error];
                    }
                  }];
  });
}

- (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 != NSURLErrorCancelled && error.code != 57) {
        [strongSelf notifyDisconnect:error];
        [strongSelf disconnectInternal];
      }
      return;
    }

    if (message.type == NSURLSessionWebSocketMessageTypeString) {
      NSLog(@"[DeepgramWebSocketClient] Received text: %@", message.string);
      [strongSelf handleTextMessage:message.string];
    } else if (message.type == NSURLSessionWebSocketMessageTypeData) {
      NSLog(@"[DeepgramWebSocketClient] Received binary: %lu bytes",
            (unsigned long)message.data.length);
      [strongSelf handleBinaryMessage:message.data];
    }

    [strongSelf receiveMessage];
  }];
}

- (void)handleTextMessage:(NSString *)text {
  if (text.length == 0) {
    return;
  }

  NSData *data = [text dataUsingEncoding:NSUTF8StringEncoding];
  if (!data) {
    return;
  }

  NSError *jsonError = nil;
  NSDictionary *json = [NSJSONSerialization JSONObjectWithData:data
                                                       options:0
                                                         error:&jsonError];
  if (jsonError) {
    [self reportError:jsonError];
    return;
  }

  NSString *errorMessage = json[@"error"];
  if (errorMessage.length > 0) {
    [self reportErrorWithMessage:errorMessage];
    return;
  }

  NSDictionary *channel = json[@"channel"];
  if (![channel isKindOfClass:[NSDictionary class]]) {
    return;
  }

  NSArray *alternatives = channel[@"alternatives"];
  if (![alternatives isKindOfClass:[NSArray class]] ||
      alternatives.count == 0) {
    return;
  }

  NSDictionary *firstAlt = alternatives.firstObject;
  NSString *transcript = firstAlt[@"transcript"] ?: @"";
  BOOL isFinal =
      [json[@"is_final"] boolValue] || [json[@"speech_final"] boolValue];

  if (transcript.length == 0) {
    return;
  }

  dispatch_async(dispatch_get_main_queue(), ^{
    if (isFinal) {
      if ([self.delegate respondsToSelector:@selector
                         (deepgramClientDidReceiveFinalTranscript:)]) {
        [self.delegate deepgramClientDidReceiveFinalTranscript:transcript];
      }
    } else {
      if ([self.delegate respondsToSelector:@selector
                         (deepgramClientDidReceiveInterimTranscript:)]) {
        [self.delegate deepgramClientDidReceiveInterimTranscript:transcript];
      }
    }
  });
}

- (void)handleBinaryMessage:(NSData *)data {
}

- (void)disconnectInternal {
  self.connected = NO;
  self.audioSendingEnabled = NO;

  if (self.webSocketTask) {
    [self.webSocketTask
        cancelWithCloseCode:NSURLSessionWebSocketCloseCodeNormalClosure
                     reason:nil];
    self.webSocketTask = nil;
  }

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

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

- (void)reportErrorWithMessage:(NSString *)message {
  NSError *error =
      [NSError errorWithDomain:kDeepgramWebSocketClientErrorDomain
                          code:-1
                      userInfo:@{NSLocalizedDescriptionKey : message ?: @""}];
  [self reportError:error];
}

- (void)notifyDisconnect:(NSError *_Nullable)error {
  self.connected = NO;

  dispatch_async(dispatch_get_main_queue(), ^{
    if ([self.delegate
            respondsToSelector:@selector(deepgramClientDidDisconnect:)]) {
      [self.delegate deepgramClientDidDisconnect:error];
    }
  });
}

#pragma mark - NSURLSessionWebSocketDelegate

- (void)URLSession:(NSURLSession *)session
          webSocketTask:(NSURLSessionWebSocketTask *)webSocketTask
    didOpenWithProtocol:(NSString *)protocol {
  self.connected = YES;
  NSLog(@"[DeepgramWebSocketClient] Connected");
  dispatch_async(dispatch_get_main_queue(), ^{
    if ([self.delegate
            respondsToSelector:@selector(deepgramClientDidConnect)]) {
      [self.delegate deepgramClientDidConnect];
    }
  });
}

- (void)URLSession:(NSURLSession *)session
       webSocketTask:(NSURLSessionWebSocketTask *)webSocketTask
    didCloseWithCode:(NSURLSessionWebSocketCloseCode)closeCode
              reason:(NSData *)reason {
  if (!self.webSocketTask) {
    return;
  }
  NSLog(@"[DeepgramWebSocketClient] Closed with code: %ld", (long)closeCode);
  [self notifyDisconnect:nil];
  [self disconnectInternal];
}

@end