//
//  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 ! autovideosink sync=false
 */

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

const NSUInteger FRAME_LENGTH_LEN = 2;
const NSUInteger MAX_RTP_PACKET_SIZE = 65535;   // because length field is two bytes (RFC 4571)
const NSUInteger RTP_HEADER_LEN = 12;
const UInt8 DEFAULT_PAYLOAD_TYPE = 96;
const NSUInteger FU_INDICATOR_LEN = 1;
const NSUInteger FU_HEADER_LEN = 1;
const UInt8 TYPE_FU_A = 0x1C;
const NSUInteger CLOCK_RATE = 90000;    // we use 90 kHz clock rate ([5.1] in RFC 6184)

// write 2-byte value in network byte order
static inline void writeShortInNBO(UInt8 *p, UInt16 value) {
    p[0] = (value >> 8) & 0xFF;
    p[1] = value & 0xFF;
}

// write 4-byte value in network byte order
static inline void writeLongInNBO(UInt8 *p, UInt32 value) {
    p[0] = (value >> 24) & 0xFF;
    p[1] = (value >> 16) & 0xFF;
    p[2] = (value >> 8) & 0xFF;
    p[3] = value & 0xFF;
}

NS_ASSUME_NONNULL_BEGIN

@interface SDLRTPH264Packetizer ()
@property (assign, nonatomic) UInt32 initialTimestamp;
@property (assign, nonatomic) UInt16 sequenceNum;
@end

@implementation SDLRTPH264Packetizer

- (instancetype)init {
    if (self = [super init]) {
        _payloadType = DEFAULT_PAYLOAD_TYPE;
        // 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);
        _ssrc = arc4random_uniform(UINT32_MAX);
    }

    return self;
}

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

- (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];
        NSUInteger nalUnitLength = [nalUnit length];
        BOOL isLast = (i + 1) == nalUnitsCount;

        if (RTP_HEADER_LEN + nalUnitLength > MAX_RTP_PACKET_SIZE) {
            // Split into multiple Fragmentation Units ([5.8] in RFC 6184)
            UInt8 firstByte;
            [nalUnit getBytes:&firstByte length:1];
            BOOL isFirstFragment = YES;
            BOOL isLastFragment = NO;
            NSUInteger offset = 1;  // we have already read the first byte

            while (offset < nalUnitLength) {
                NSUInteger payloadLength = MAX_RTP_PACKET_SIZE - (RTP_HEADER_LEN + FU_INDICATOR_LEN + FU_HEADER_LEN);
                if (nalUnitLength - offset <= payloadLength) {
                    payloadLength = nalUnitLength - offset;
                    isLastFragment = YES;
                }
                NSUInteger packetSize = RTP_HEADER_LEN + FU_INDICATOR_LEN + FU_HEADER_LEN + payloadLength;
                NSUInteger frameSize = FRAME_LENGTH_LEN + packetSize;

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

                p += [self writeFrameHeader:p packetSize:packetSize];
                p += [self writeRTPHeader:p marker:isLast pts:presentationTimestamp];

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

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

                isFirstFragment = NO;
            }
        } else {
            // Use Single NAL Unit Packet ([5.6] in RFC 6184)
            NSUInteger packetSize = RTP_HEADER_LEN + nalUnitLength;
            NSUInteger frameSize = FRAME_LENGTH_LEN + packetSize;

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

            p += [self writeFrameHeader:p packetSize:packetSize];
            p += [self writeRTPHeader:p marker:isLast pts:presentationTimestamp];
            [nalUnit getBytes:p length:nalUnitLength];

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

    return rtpFrames;
}

/**
 * Write RTP Frame header (length field) to supplied buffer.
 *
 * @param p 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)writeFrameHeader:(UInt8 *)p packetSize:(NSUInteger)packetSize {
    NSAssert(packetSize <= MAX_RTP_PACKET_SIZE, @"RTP packet is too big");
    writeShortInNBO(p, (UInt16)packetSize);
    return FRAME_LENGTH_LEN;
}

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

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

    self.sequenceNum++;
    return RTP_HEADER_LEN;
}

@end

NS_ASSUME_NONNULL_END