//
//  FlutterSocketConnectionFrameReader.m
//  RCTWebRTC
//
//  Created by Alex-Dan Bumbu on 06/01/2021.
//

#include <mach/mach_time.h>

#import <ReplayKit/ReplayKit.h>
#import <WebRTC/RTCCVPixelBuffer.h>
#import <WebRTC/RTCVideoFrameBuffer.h>

#import "FlutterSocketConnection.h"
#import "FlutterSocketConnectionFrameReader.h"

const NSUInteger kMaxReadLength = 10 * 1024;

@interface Message : NSObject

@property(nonatomic, assign, readonly) CVImageBufferRef imageBuffer;
@property(nonatomic, copy, nullable) void (^didComplete)(BOOL succes, Message* message);

- (NSInteger)appendBytes:(UInt8*)buffer length:(NSUInteger)length;

@end

@interface Message ()

@property(nonatomic, assign) CVImageBufferRef imageBuffer;
@property(nonatomic, assign) int imageOrientation;
@property(nonatomic, assign) CFHTTPMessageRef framedMessage;

@end

@implementation Message

- (instancetype)init {
  self = [super init];
  if (self) {
    self.imageBuffer = NULL;
  }

  return self;
}

- (void)dealloc {
  CVPixelBufferRelease(_imageBuffer);
}

/** Returns the amount of missing bytes to complete the message, or -1 when not enough bytes were
 * provided to compute the message length */
- (NSInteger)appendBytes:(UInt8*)buffer length:(NSUInteger)length {
  if (!_framedMessage) {
    _framedMessage = CFHTTPMessageCreateEmpty(kCFAllocatorDefault, false);
  }

  CFHTTPMessageAppendBytes(_framedMessage, buffer, length);
  if (!CFHTTPMessageIsHeaderComplete(_framedMessage)) {
    return -1;
  }

  NSInteger contentLength = [CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(
      _framedMessage, (__bridge CFStringRef) @"Content-Length")) integerValue];
  NSInteger bodyLength =
      (NSInteger)[CFBridgingRelease(CFHTTPMessageCopyBody(_framedMessage)) length];

  NSInteger missingBytesCount = contentLength - bodyLength;
  if (missingBytesCount == 0) {
    BOOL success = [self unwrapMessage:self.framedMessage];
    self.didComplete(success, self);

    CFRelease(self.framedMessage);
    self.framedMessage = NULL;
  }

  return missingBytesCount;
}

// MARK: Private Methods

- (CIContext*)imageContext {
  // Initializing a CIContext object is costly, so we use a singleton instead
  static CIContext* imageContext = nil;
  static dispatch_once_t onceToken;
  dispatch_once(&onceToken, ^{
    imageContext = [[CIContext alloc] initWithOptions:nil];
  });

  return imageContext;
}

- (BOOL)unwrapMessage:(CFHTTPMessageRef)framedMessage {
  size_t width = [CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(
      _framedMessage, (__bridge CFStringRef) @"Buffer-Width")) integerValue];
  size_t height = [CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(
      _framedMessage, (__bridge CFStringRef) @"Buffer-Height")) integerValue];
  _imageOrientation = [CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(
      _framedMessage, (__bridge CFStringRef) @"Buffer-Orientation")) intValue];

  NSData* messageData = CFBridgingRelease(CFHTTPMessageCopyBody(_framedMessage));

  // Copy the pixel buffer
  CVReturn status = CVPixelBufferCreate(kCFAllocatorDefault, width, height,
                                        kCVPixelFormatType_32BGRA, NULL, &_imageBuffer);
  if (status != kCVReturnSuccess) {
    NSLog(@"CVPixelBufferCreate failed");
    return false;
  }

  [self copyImageData:messageData toPixelBuffer:&_imageBuffer];

  return true;
}

- (void)copyImageData:(NSData*)data toPixelBuffer:(CVPixelBufferRef*)pixelBuffer {
  CVPixelBufferLockBaseAddress(*pixelBuffer, 0);

  CIImage* image = [CIImage imageWithData:data];
  [self.imageContext render:image toCVPixelBuffer:*pixelBuffer];

  CVPixelBufferUnlockBaseAddress(*pixelBuffer, 0);
}

@end

// MARK: -

@interface FlutterSocketConnectionFrameReader () <NSStreamDelegate>

@property(nonatomic, strong) FlutterSocketConnection* connection;
@property(nonatomic, strong) Message* message;

@end

@implementation FlutterSocketConnectionFrameReader {
  mach_timebase_info_data_t _timebaseInfo;
  NSInteger _readLength;
  int64_t _startTimeStampNs;
}

- (instancetype)initWithDelegate:(__weak id<RTCVideoCapturerDelegate>)delegate {
  self = [super initWithDelegate:delegate];
  if (self) {
    mach_timebase_info(&_timebaseInfo);
  }

  return self;
}

- (void)startCaptureWithConnection:(FlutterSocketConnection*)connection {
  _startTimeStampNs = -1;

  self.connection = connection;
  self.message = nil;

  [self.connection openWithStreamDelegate:self];
}

- (void)stopCapture {
  [self.connection close];
}

// MARK: Private Methods

- (void)readBytesFromStream:(NSInputStream*)stream {
  if (!stream.hasBytesAvailable) {
    return;
  }

  if (!self.message) {
    self.message = [[Message alloc] init];
    _readLength = kMaxReadLength;

    __weak __typeof__(self) weakSelf = self;
    self.message.didComplete = ^(BOOL success, Message* message) {
      if (success) {
        [weakSelf didCaptureVideoFrame:message.imageBuffer
                       withOrientation:message.imageOrientation];
      }

      weakSelf.message = nil;
    };
  }

  uint8_t buffer[_readLength];
  NSInteger numberOfBytesRead = [stream read:buffer maxLength:_readLength];
  if (numberOfBytesRead < 0) {
    NSLog(@"error reading bytes from stream");
    return;
  }

  _readLength = [self.message appendBytes:buffer length:numberOfBytesRead];
  if (_readLength == -1 || _readLength > kMaxReadLength) {
    _readLength = kMaxReadLength;
  }
}

- (void)didCaptureVideoFrame:(CVPixelBufferRef)pixelBuffer
             withOrientation:(CGImagePropertyOrientation)orientation {
  int64_t currentTime = mach_absolute_time();
  int64_t currentTimeStampNs = currentTime * _timebaseInfo.numer / _timebaseInfo.denom;

  if (_startTimeStampNs < 0) {
    _startTimeStampNs = currentTimeStampNs;
  }

  RTCCVPixelBuffer* rtcPixelBuffer = [[RTCCVPixelBuffer alloc] initWithPixelBuffer:pixelBuffer];
  int64_t frameTimeStampNs = currentTimeStampNs - _startTimeStampNs;

  RTCVideoRotation rotation;
  switch (orientation) {
    case kCGImagePropertyOrientationLeft:
      rotation = RTCVideoRotation_90;
      break;
    case kCGImagePropertyOrientationDown:
      rotation = RTCVideoRotation_180;
      break;
    case kCGImagePropertyOrientationRight:
      rotation = RTCVideoRotation_270;
      break;
    default:
      rotation = RTCVideoRotation_0;
      break;
  }

  RTCVideoFrame* videoFrame = [[RTCVideoFrame alloc] initWithBuffer:[rtcPixelBuffer toI420]
                                                           rotation:rotation
                                                        timeStampNs:frameTimeStampNs];

  [self.delegate capturer:self didCaptureVideoFrame:videoFrame];
}

@end

@implementation FlutterSocketConnectionFrameReader (NSStreamDelegate)

- (void)stream:(NSStream*)aStream handleEvent:(NSStreamEvent)eventCode {
  switch (eventCode) {
    case NSStreamEventOpenCompleted:
      NSLog(@"server stream open completed");
      break;
    case NSStreamEventHasBytesAvailable:
      [self readBytesFromStream:(NSInputStream*)aStream];
      break;
    case NSStreamEventEndEncountered:
      NSLog(@"server stream end encountered");
      [self stopCapture];
      break;
    case NSStreamEventErrorOccurred:
      NSLog(@"server stream error encountered: %@", aStream.streamError.localizedDescription);
      break;

    default:
      break;
  }
}

@end
