//
//  SDLRTPH264Packetizer.m
//  SmartDeviceLink-iOS
//
//  Created by Sho Amano on 4/11/17.
//  Copyright © 2017 Xevo Inc. All rights reserved.
//

/*
 * Note for testing.
 * The RTP stream generated by this packetizer can be tested with GStreamer (1.4 or later).
 * Assuming that "VideoStreamPort" is configured as 5050 in smartDeviceLink.ini, here is the
 * GStreamer pipeline that receives the stream, decode it and render it:
 *
 * $ gst-launch-1.0 souphttpsrc location=http://127.0.0.1:5050 ! "application/x-rtp-stream" ! rtpstreamdepay ! "application/x-rtp,media=(string)video,clock-rate=90000,encoding-name=(string)H264" ! rtph264depay ! "video/x-h264, stream-format=(string)avc, alignment=(string)au" ! avdec_h264 ! videoconvert ! ximagesink sync=false
 */

#import <Foundation/Foundation.h>
#import "SDLRTPH264Packetizer.h"
#import "SDLLogMacros.h"

const NSUInteger FrameLengthLen = 2;
const NSUInteger MaxRTPPacketSize = 65535;  // because length field is two bytes (RFC 4571)
const NSUInteger RTPHeaderLen = 12;
const UInt8 DefaultPayloadType = 96;
const NSUInteger FragmentationUnitIndicatorLen = 1;
const NSUInteger FragmentationUnitHeaderLen = 1;
const UInt8 FragmentationUnitVersionA = 0x1C;
const NSUInteger ClockRate = 90000; // we use 90 kHz clock rate ([5.1] in RFC 6184)

/**
 * write 2-byte value into buffer in network byte order
 *
 * @param buffer buffer to write the value
 * @param value value to write
 */
static inline void sdl_writeShortInNetworkByteOrder(UInt8 *buffer, UInt16 value) {
    buffer[0] = (value >> 8) & 0xFF;
    buffer[1] = value & 0xFF;
}

/**
 * write 4-byte value into buffer in network byte order
 *
 * @param buffer buffer to write the value
 * @param value value to write
 */
static inline void sdl_writeLongInNetworkByteOrder(UInt8 *buffer, UInt32 value) {
    buffer[0] = (value >> 24) & 0xFF;
    buffer[1] = (value >> 16) & 0xFF;
    buffer[2] = (value >> 8) & 0xFF;
    buffer[3] = value & 0xFF;
}

NS_ASSUME_NONNULL_BEGIN

@interface SDLRTPH264Packetizer ()

@property (assign, nonatomic) UInt32 initialTimestamp;
@property (assign, nonatomic) UInt16 sequenceNum;
@property (assign, nonatomic) UInt32 ssrc;

@end

@implementation SDLRTPH264Packetizer

- (instancetype)initWithSSRC:(UInt32)ssrc {
    self = [self init];
    _ssrc = ssrc;

    return self;
}

- (instancetype)init {
    self = [super init];
    if (!self) {
        return nil;
    }

    _payloadType = DefaultPayloadType;
    // initial value of the sequence number and timestamp should be random ([5.1] in RFC3550)
    _initialTimestamp = arc4random_uniform(UINT32_MAX);
    _sequenceNum = (UInt16)arc4random_uniform(UINT16_MAX);

    return self;
}

- (void)setPayloadType:(UInt8)payloadType {
    if (payloadType <= 127) {
        _payloadType = payloadType;
    } else {
        _payloadType = DefaultPayloadType;
    }
}

- (nullable NSArray<NSData *> *)createPackets:(NSArray<NSData *> *)nalUnits
                        presentationTimestamp:(double)presentationTimestamp {
    NSMutableArray *rtpFrames = [NSMutableArray array];
    NSUInteger nalUnitsCount = nalUnits.count;

    for (NSUInteger i = 0; i < nalUnitsCount; i++) {
        NSData *nalUnit = nalUnits[i];
        BOOL isLast = ((i + 1) == nalUnitsCount);

        if (RTPHeaderLen + nalUnit.length > MaxRTPPacketSize) {
            // Split into multiple Fragmentation Units ([5.8] in RFC 6184)
            if (![self sdl_addRTPFramesWithFragmentationUnits:rtpFrames nalUnit:nalUnit presentationTimestamp:presentationTimestamp isLast:isLast]) {
                return nil;
            }
        } else {
            // Use Single NAL Unit Packet ([5.6] in RFC 6184)
            if (![self sdl_addRTPFrameWithSingleNALUnit:rtpFrames nalUnit:nalUnit presentationTimestamp:presentationTimestamp isLast:isLast]) {
                return nil;
            }
        }
    }

    return [rtpFrames copy];
}

/**
 * Create and add a RTP frame from a NAL unit without splitting
 *
 * @param rtpFrames the array to which created RTP frame is added
 * @param nalUnit NAL unit to create RTP frame from
 * @param presentationTimestamp presentation timestamp in seconds
 * @param isLast mark YES if this is the last NAL unit of a video frame
 *
 * @return YES if successful, NO if memory error occurred
 */
- (BOOL)sdl_addRTPFrameWithSingleNALUnit:(NSMutableArray<NSData *> *)rtpFrames
                                 nalUnit:(NSData *)nalUnit
                   presentationTimestamp:(double)presentationTimestamp
                                  isLast:(BOOL)isLast {
    NSUInteger nalUnitLength = nalUnit.length;
    NSUInteger packetSize = RTPHeaderLen + nalUnitLength;
    NSUInteger frameSize = FrameLengthLen + packetSize;

    NSAssert(RTPHeaderLen + nalUnitLength <= MaxRTPPacketSize, @"This NAL unit doesn't fit into single RTP packet");

    UInt8 *buffer = malloc(frameSize);
    if (buffer == NULL) {
        SDLLogE(@"malloc() error");
        return NO;
    }
    UInt8 *writePointer = buffer;

    writePointer += [self sdl_writeFrameHeader:writePointer packetSize:packetSize];
    writePointer += [self sdl_writeRTPHeader:writePointer marker:isLast presentationTimestamp:presentationTimestamp];
    [nalUnit getBytes:writePointer length:nalUnitLength];

    NSData *rtpFrame = [NSData dataWithBytesNoCopy:buffer length:frameSize];
    [rtpFrames addObject:rtpFrame];

    return YES;
}

/**
 * Create and add a RTP frames by splitting a NAL unit into multiple Fragmentation Units
 *
 * @param rtpFrames the array to which created RTP frames are added
 * @param nalUnit NAL unit to create RTP frames from
 * @param presentationTimestamp presentation timestamp in seconds
 * @param isLast mark YES if this is the last NAL unit of a video frame
 *
 * @return YES if successful, NO if memory error occurred
 */
- (BOOL)sdl_addRTPFramesWithFragmentationUnits:(NSMutableArray<NSData *> *)rtpFrames
                                       nalUnit:(NSData *)nalUnit
                         presentationTimestamp:(double)presentationTimestamp
                                        isLast:(BOOL)isLast {
    UInt8 firstByte;
    [nalUnit getBytes:&firstByte length:1];
    BOOL isFirstFragment = YES;
    BOOL isLastFragment = NO;
    NSUInteger nalUnitLength = nalUnit.length;
    NSUInteger offset = 1;  // we have already read the first byte

    while (offset < nalUnitLength) {
        NSUInteger payloadLength = MaxRTPPacketSize - (RTPHeaderLen + FragmentationUnitIndicatorLen + FragmentationUnitHeaderLen);
        if (nalUnitLength - offset <= payloadLength) {
            payloadLength = nalUnitLength - offset;
            isLastFragment = YES;
        }
        NSUInteger packetSize = RTPHeaderLen + FragmentationUnitIndicatorLen + FragmentationUnitHeaderLen + payloadLength;
        NSUInteger frameSize = FrameLengthLen + packetSize;

        UInt8 *buffer = malloc(frameSize);
        if (buffer == NULL) {
            SDLLogE(@"malloc() error");
            return NO;
        }
        UInt8 *writePointer = buffer;

        writePointer += [self sdl_writeFrameHeader:writePointer packetSize:packetSize];
        writePointer += [self sdl_writeRTPHeader:writePointer marker:isLast presentationTimestamp:presentationTimestamp];

        // Fragmentation Unit indicator
        *writePointer++ = (firstByte & 0xE0) | FragmentationUnitVersionA;
        // Fragmentation Unit header
        *writePointer++ = (isFirstFragment ? 0x80 : isLastFragment ? 0x40 : 0) | (firstByte & 0x1F);
        // Fragmentation Unit payload
        [nalUnit getBytes:writePointer range:NSMakeRange(offset, payloadLength)];
        offset += payloadLength;

        NSData *rtpFrame = [NSData dataWithBytesNoCopy:buffer length:frameSize];
        [rtpFrames addObject:rtpFrame];

        isFirstFragment = NO;
    }

    return YES;
}

/**
 * Write RTP Frame header (length field) to supplied buffer.
 *
 * @param frameHeaderBuffer the buffer in which a header is written.
 * @param packetSize size of a RTP packet that follows to this frame header.
 *
 * @return number of bytes written, which is always 2.
 */
- (NSUInteger)sdl_writeFrameHeader:(UInt8 *)frameHeaderBuffer packetSize:(NSUInteger)packetSize {
    NSAssert(packetSize <= MaxRTPPacketSize, @"RTP packet is too big");
    sdl_writeShortInNetworkByteOrder(frameHeaderBuffer, (UInt16)packetSize);
    return FrameLengthLen;
}

/**
 * Write RTP header to supplied buffer.
 *
 * @param rtpHeaderBuffer the buffer in which a header is written.
 * @param isLast whether this is the last packet of an access unit.
 * @param presentationTimestamp presentation timestamp in seconds.
 *
 * @return number of bytes written, which is always 12.
 */
- (NSUInteger)sdl_writeRTPHeader:(UInt8 *)rtpHeaderBuffer marker:(BOOL)isLast presentationTimestamp:(double)presentationTimestamp {
    UInt32 presentationTimestampIn90kHz = (UInt32)(presentationTimestamp * ClockRate);

    // Version = 2, Padding = 0, Extension = 0, CSRC count = 0
    rtpHeaderBuffer[0] = 0x80;
    // Marker = isLast, Payload type = self.payloadType
    rtpHeaderBuffer[1] = (isLast ? 0x80 : 0) | (self.payloadType & 0x7F);
    sdl_writeShortInNetworkByteOrder(rtpHeaderBuffer + 2, self.sequenceNum);
    sdl_writeLongInNetworkByteOrder(rtpHeaderBuffer + 4, self.initialTimestamp + presentationTimestampIn90kHz);
    sdl_writeLongInNetworkByteOrder(rtpHeaderBuffer + 8, self.ssrc);

    self.sequenceNum++;
    return RTPHeaderLen;
}

@end

NS_ASSUME_NONNULL_END