#include "RtmpPusher.h"

/**
 * 初始化winsock
 *
 * @成功则返回1 , 失败则返回相应错误代码
 */
int InitSockets()
{
#ifdef WIN32
    WORD version;
    WSADATA wsaData;
    version = MAKEWORD(1, 1);
    return (WSAStartup(version, &wsaData) == 0);
#else
    return TRUE;
#endif
}

bool RtmpPusher::IfConnect()
{
     return mIfConnected;
}

int RtmpPusher::RTMP264_Connect(const char* url)
{
    InitSockets();
    m_pRtmp = RTMP_Alloc();
    RTMP_Init(m_pRtmp);
    /*设置URL*/
    if (RTMP_SetupURL(m_pRtmp, (char*)url) == FALSE)
    {
        RTMP_Free(m_pRtmp);
        return -1;
    }
    /*设置可写,即发布流,这个函数必须在连接前使用,否则无效*/
    RTMP_EnableWrite(m_pRtmp);
    /*连接服务器*/
    if (RTMP_Connect(m_pRtmp, NULL) == FALSE)
    {
        RTMP_Free(m_pRtmp);
        return -1;
    }

    /*连接流*/
    if (RTMP_ConnectStream(m_pRtmp, 0) == FALSE)
    {
        RTMP_Close(m_pRtmp);
        RTMP_Free(m_pRtmp);
        return -1;
    }

    this->mUrl = string(url);
    this->mIfConnected = true;
    return 0;
}
/**
 * 释放winsock
 *
 * @成功则返回0 , 失败则返回相应错误代码
 */
inline void CleanupSockets()
{
#ifdef WIN32
    WSACleanup();
#endif
}

void RtmpPusher::RTMP264_Close()
{
    mMux.lock();
    if (m_pRtmp)
    {
        RTMP_Close(m_pRtmp);
        RTMP_Free(m_pRtmp);
        m_pRtmp = NULL;
    }
    mMux.unlock();
    CleanupSockets();
}
RTMPPacket* gPacket = nullptr;

int RtmpPusher::SendPacket(unsigned int nPacketType, unsigned char * data,
    unsigned int size, unsigned int nTimestamp)
{
    static bool once = true;
    /*分配包内存和初始化,len为包体长度*/
    if(nullptr == gPacket)
        gPacket = (RTMPPacket *)malloc(640*720*3 + size);
    memset(gPacket, 0, RTMP_HEAD_SIZE);
    /*包体内存*/

    gPacket->m_body = (char *)gPacket + RTMP_HEAD_SIZE;
    gPacket->m_nBodySize = size;
    memcpy(gPacket->m_body, data, size);

    gPacket->m_hasAbsTimestamp = 0;
    gPacket->m_packetType = nPacketType; /*此处为类型有两种一种是音频,一种是视频*/
    gPacket->m_nInfoField2 = m_pRtmp->m_stream_id;
    gPacket->m_nChannel = 0x04;


    gPacket->m_headerType = RTMP_PACKET_SIZE_LARGE;
    if (RTMP_PACKET_TYPE_AUDIO == nPacketType && size != 4)
    {
        gPacket->m_headerType = RTMP_PACKET_SIZE_MEDIUM;
    }

    gPacket->m_nTimeStamp = nTimestamp;
    /*发送*/
    int nRet = 0;
    if (RTMP_IsConnected(m_pRtmp))
    {
        nRet = RTMP_SendPacket(m_pRtmp, gPacket, FALSE); /*TRUE为放进发送队列,FALSE是不放进发送队列,直接发送*/
    }
    else {
        if (once) {
            once = false;
        }
    }
    /*释放内存*/
    //free(gPacket);
    return nRet;
}

int RtmpPusher::SendVideoPacket(unsigned int nPacketType,
    unsigned char * data, unsigned int size, unsigned int nTimestamp)
{

    RTMPPacket* packet;
    /*分配包内存和初始化,len为包体长度*/
    packet = (RTMPPacket *)malloc(RTMP_HEAD_SIZE + size);
    memset(packet, 0, RTMP_HEAD_SIZE);
    /*包体内存*/
    packet->m_body = (char *)packet + RTMP_HEAD_SIZE;
    packet->m_nBodySize = size;
    memcpy(packet->m_body, data, size);
    packet->m_hasAbsTimestamp = 0;
    packet->m_packetType = nPacketType; /*此处为类型有两种一种是音频,一种是视频*/
    packet->m_nInfoField2 = m_pRtmp->m_stream_id;
    packet->m_nChannel = 0x04;

    packet->m_headerType = RTMP_PACKET_SIZE_LARGE;
    if (RTMP_PACKET_TYPE_AUDIO == nPacketType && size != 4)
    {
        packet->m_headerType = RTMP_PACKET_SIZE_MEDIUM;
    }
    packet->m_nTimeStamp = nTimestamp;
    /*发送*/
    int nRet = 0;
    if (RTMP_IsConnected(m_pRtmp))
    {
        nRet = RTMP_SendPacket(m_pRtmp, packet, TRUE); /*TRUE为放进发送队列,FALSE是不放进发送队列,直接发送*/
    }
    /*释放内存*/
    free(packet);
    return 0;
}

RtmpPusher::RtmpPusher()
    :mThread(nullptr),
    mIfConnected(false)
{

}


RtmpPusher::~RtmpPusher()
{
    if (m_pRtmp)
    {
        RTMP_Close(m_pRtmp);
        RTMP_Free(m_pRtmp);
        m_pRtmp = NULL;
    }
    CleanupSockets();
}

void H264RtmpPuser::OnAudioEncode(const void * frameaddress, uint32_t framelen,uint16_t pts)
{
    uint8_t *pack = (uint8_t*)malloc(framelen);
    memcpy(pack, frameaddress, framelen);

    mMux.lock();
    Buffer buf;
    buf.buf = (uint8_t *)pack;
    buf.len = framelen;
    buf.type = PAYLOAD_TYPE_AUDIO;
    this->mPack.push(buf);
    mMux.unlock();
    this->mAudioPts = pts;
}

H264RtmpPuser::H264RtmpPuser()
{
    this->metaData.Pps = nullptr;
    this->metaData.Sps = nullptr;
    this->metaData.nPpsLen = 0;
    this->metaData.nSpsLen = 0;
    mFirtACC = false;
}

int H264RtmpPuser::sortAndSendNal(uint8_t * data, int len)
{
    int i = 0;
    uint8_t * nalhead = nullptr;
    uint8_t * naltail = nullptr;
    uint32_t size = 0;

    static int timestamp = 0;
    timestamp += 1000 / 25;
    if (nullptr == data) {
        return -1;
    }
    while (i < len)
    {
        // sps pps p frame
        if ((data[i] == 0x00) && (data[i + 1] == 0x00)
            && ((data[i + 2] == 0x00) && (data[i + 3] == 0x01) || (data[i + 2] == 0x01))) {
            if ((nalhead == nullptr) && (i == 0) ) {
                if ((data[i + 3] == 0x01) && (data[i + 4] == 0x41)) {  //p 帧直接发
                    nalhead = data;
                    naltail = data + (len);
                    size = naltail - nalhead;
                    this->SendH264Packet(nalhead, size, 0, timestamp);
                    return 0;
                }
                //sps 帧进行解包
                if ((data[i + 3] == 0x01) && (data[i + 4] == 0x67)) {  // sps or pps or sei
                    nalhead = data;
                    i += 1;
                }
                //sei
                if ((data[i + 2] == 0x01) && (data[i + 3] == 0x06)) {
                    i += 1;
                }
            }
            else {
                // i frame
                if ((data[i + 2] == 0x01) && (data[i + 3] == 0x65)) {
                    naltail = data + i;
                    size = naltail - nalhead;
                    this->SendH264Packet(nalhead, size, 0, timestamp);
                    nalhead = data + i;
                    naltail = data + (len);
                    size = naltail - nalhead;
                    this->SendH264Packet(nalhead, size, 0, timestamp);
                    return 0;
                }
                //pps
                if ((data[i + 3] == 0x01) && (data[i + 4] == 0x68)) {  // sps or pps or sei
                    naltail = data + i;
                    size = naltail - nalhead;
                    this->SendH264Packet(nalhead, size, 0, timestamp);
                    nalhead = data + i;
                    i += 3;
                }//sps
                if ((data[i + 3] == 0x01) && (data[i + 4] == 0x67)) {  // sps or pps or sei
                    nalhead = data + i;
                    i += 3;
                }
                //sei
                if ((data[i + 3] == 0x01) && (data[i + 4] == 0x06)) {  // sps or pps or sei
                    naltail = data + i;
                    size = naltail - nalhead;
                    this->SendH264Packet(nalhead, size, 0, timestamp);
                    nalhead = data + i;
                    i += 3;
                }
                // sps pps or sei
            }
            // 跳过00 00 00 00 01的情况
        }
        i++;
    }
    return 0;
}


// 视频同步包，详细结构请见https://blog.csdn.net/liwf616/article/details/51596373
int H264RtmpPuser::SendVideoSpsPps(unsigned char * pps,
    int pps_len, unsigned char * sps,
    int sps_len,unsigned int nTimeStamp)
{
    RTMPPacket * packet = NULL;//rtmp包结构
    unsigned char * body = NULL;
    int i;
    packet = (RTMPPacket *)malloc(RTMP_HEAD_SIZE + 1024);
    //RTMPPacket_Reset(packet);//重置packet状态
    memset(packet, 0, RTMP_HEAD_SIZE + 1024);
    packet->m_body = (char *)packet + RTMP_HEAD_SIZE;
    body = (unsigned char *)packet->m_body;
    i = 0;
    // FrameType == 1，CodecID == 7，
    body[i++] = 0x17;
    //AVCPacketType
    body[i++] = 0x00;

    //CompositionTime
    body[i++] = 0x00;
    body[i++] = 0x00;
    body[i++] = 0x00;

    /*AVCDecoderConfigurationRecord*/
    body[i++] = 0x01;
    body[i++] = sps[1];
    body[i++] = sps[2];
    body[i++] = sps[3];
    body[i++] = 0xff;

    /*sps*/
    body[i++] = 0xe1;
    body[i++] = (sps_len >> 8) & 0xff;
    body[i++] = sps_len & 0xff;
    memcpy(&body[i], sps, sps_len);
    i += sps_len;

    /*pps*/
    body[i++] = 0x01;
    body[i++] = (pps_len >> 8) & 0xff;
    body[i++] = (pps_len) & 0xff;
    memcpy(&body[i], pps, pps_len);
    i += pps_len;

    packet->m_packetType = RTMP_PACKET_TYPE_VIDEO;
    packet->m_nBodySize = i;
    packet->m_nChannel = 0x04;
    packet->m_nTimeStamp = nTimeStamp;
    packet->m_hasAbsTimestamp = 0;
    packet->m_headerType = RTMP_PACKET_SIZE_MEDIUM;
    packet->m_nInfoField2 = m_pRtmp->m_stream_id;

    /*调用发送接口*/
    int nRet = RTMP_SendPacket(m_pRtmp, packet, TRUE);
    free(packet);    //释放内存
    return nRet;
}

int H264RtmpPuser::SendAudioData(unsigned char * dat,
    unsigned int size, unsigned int nTimeStamp)
{
    return 0;
}

int H264RtmpPuser::SendH264Packet(unsigned char * data,
    unsigned int size, int bIsKeyFrame, unsigned int nTimeStamp)
{
    if(data == NULL){
        return false;
    }
    unsigned int nal_type = 0;
    // 小帧应该是PPS或者SPS
    if ((data[0] != 0x00) || (data[1] != 0x00)
        || ((data[2] != 0x00)&&data[2]!= 0x01)) {
        return false;
    }
    //Debuger::Debug(L"%02x %02x %02x %02x %02x %02d\r\n",
    //	data[0],data[1],data[2],data[3],data[4],size);
    if (data[2] == 0x01) {
        nal_type = data[3];
    }
    if (data[3] == 0x01) {
        nal_type = data[4];
    }
    switch (nal_type)
    {
    case 0x67:  //just update sps and pps
        if (NULL == metaData.Sps)
            metaData.Sps = (unsigned char *)malloc(size - 4);
        h264_decode_sps(data + 4, size - 4, metaData.nWidth, metaData.nHeight, metaData.nFrameRate);
        metaData.nSpsLen = size - 4;
        memcpy(this->metaData.Sps, data + 4, size - 4);
        break;
    case 0x68: //just update sps and pps
        this->metaData.nPpsLen = size - 4;
        if (NULL == metaData.Pps) metaData.Pps = (unsigned char *)malloc(size - 4);
        memcpy(this->metaData.Pps, data + 4, size - 4);
        break;
    case 0x41:  //p frame
        this->sendDataPackH264(data + 4, size - 4, 0, nTimeStamp);
        break;
    case 0x65: //i frame
        this->sendDataPackH264(data + 3, size - 3, 1, nTimeStamp);
        break;
    case 0x06:
        size = size;
        //this->sendDataPack(data + 4, size - 4, 0, nTimeStamp);
        break;
    default:
        break;
    }

}
unsigned char *gBody = nullptr;
int H264RtmpPuser::sendDataPackH264(unsigned char * data,
    unsigned int size, int bIsKeyFrame, unsigned int nTimeStamp)
{
    if (gBody == nullptr) {
        gBody = new  unsigned char[640*720*3 + 9];
    }
    if (size < 0) {
        gBody = gBody;
    }
    memset(gBody, 0, size + 9);
    int i = 0;
    if (1 == bIsKeyFrame) {
        gBody[i++] = 0x17;// 1:Iframe  7:AVC
        gBody[i++] = 0x01;// AVC NALU
        gBody[i++] = 0x00;
        gBody[i++] = 0x00;
        gBody[i++] = 0x00;

        // NALU size
        gBody[i++] = size >> 24 & 0xff;
        gBody[i++] = size >> 16 & 0xff;
        gBody[i++] = size >> 8 & 0xff;
        gBody[i++] = size & 0xff;
        // NALU data
        memcpy(&gBody[i], data, size);
        if(metaData.Sps != nullptr)
            SendVideoSpsPps(metaData.Pps, metaData.nPpsLen, metaData.Sps,
                metaData.nSpsLen, nTimeStamp);
    }
    else {
        gBody[i++] = 0x27;// 2:Pframe  7:AVC
        gBody[i++] = 0x01;// AVC NALU
        gBody[i++] = 0x00;
        gBody[i++] = 0x00;
        gBody[i++] = 0x00;
        // NALU size
        gBody[i++] = size >> 24 & 0xff;
        gBody[i++] = size >> 16 & 0xff;
        gBody[i++] = size >> 8 & 0xff;
        gBody[i++] = size & 0xff;
        // NALU data

        memcpy(&gBody[i], data, size);
    }
    int bRet = SendPacket(RTMP_PACKET_TYPE_VIDEO, gBody, i + size, nTimeStamp);
    return bRet;
}


int H264RtmpPuser::SendAudioSync(int audioType,
    int sampleIndex, int channel, unsigned int nTimeStamp)
{
    RTMPPacket * packet = NULL;//rtmp包结构
    unsigned char * body = NULL;
    int i;

    packet = (RTMPPacket *)malloc(RTMP_HEAD_SIZE + 1024);
    //RTMPPacket_Reset(packet);//重置packet状态
    memset(packet, 0, RTMP_HEAD_SIZE + 1024);
    packet->m_body = (char *)packet + RTMP_HEAD_SIZE;
    body = (unsigned char *)packet->m_body;

    body[0] = 0xaf;
    body[1] = 0x00;

    uint16_t audioSpecConf = 0;
    audioSpecConf |= ((2 << 11) & 0xf800);  //2: AACLC
    audioSpecConf |= ((4 << 7) & 0x0780);  //4: 44khz
    audioSpecConf |= ((2 << 3) & 0x78);  //4: 2:stero
    audioSpecConf |= 0 & 0x07;			//4: 0 padding
    body[2] = (audioSpecConf >> 8) & 0xff;
    body[3] = audioSpecConf & 0xff;


    packet->m_packetType = RTMP_PACKET_TYPE_AUDIO;
    packet->m_nBodySize = 4;
    packet->m_nChannel = 0x04;
    packet->m_nTimeStamp = nTimeStamp;
    packet->m_hasAbsTimestamp = 0;
    packet->m_headerType = RTMP_PACKET_SIZE_MEDIUM;
    packet->m_nInfoField2 = m_pRtmp->m_stream_id;

    /*调用发送接口*/
    int nRet = RTMP_SendPacket(m_pRtmp, packet, TRUE);
    free(packet);    //释放内存
    return nRet;
}


int H264RtmpPuser::sendDataPackAAC(unsigned char * data,
    unsigned int size, unsigned int nTimeStamp)
{
    unsigned char *gBody = nullptr;
    static int timestamp = 0;
    timestamp += 20;
    if (!mFirtACC) {
        SendAudioSync(2,4,4, timestamp);
        mFirtACC = 1;
    }
    gBody = (unsigned char*)malloc(size + 2);
    gBody[0] = 0xAF;
    gBody[1] = 0x01;   //aac raw data
    memcpy(gBody + 2, data + 7, size - 7);
    int bRet = SendPacket(RTMP_PACKET_TYPE_AUDIO, gBody,
        size - 7 + 2, timestamp);
    free(gBody);
    return 0;
}

void H264RtmpPuser::OnGetCodeFrame(uint8_t * data, int len)
{
    static int timetamp = 0;
    timetamp += this->mTick;
    uint8_t *pack = (uint8_t*)malloc(len);
    memcpy(pack, data, len);
    mMux.lock();
    Buffer buf;
    buf.buf = pack;
    buf.len = len;
    buf.type = PAYLOAD_TYPE_VIDEO;
    this->mPack.push(buf);
    mMux.unlock();
}

void H264RtmpPuser::ProcessSend()
{
    while (this->mIfStart) {
        int len = mPack.size();
        if (!mPack.empty()) {
            mMux.lock();
            Buffer buf = mPack.front();
            mPack.pop();
            mMux.unlock();
            //如果是视频帧
            if (buf.type == PAYLOAD_TYPE_VIDEO) {
                this->sortAndSendNal(buf.buf, buf.len);

            }// 如果是音频帧
            if (buf.type == PAYLOAD_TYPE_AUDIO) {
                this->sendDataPackAAC(buf.buf, buf.len, this->mAudioPts);
            }
            free(buf.buf);
        }
        msleep(10);
    }
}

int ThreadEncode(H264RtmpPuser * p)
{
    Debuger::Debug(L"thread started\r\n");
    if (nullptr == p)
        return -1;
    p->ProcessSend();
    return 0;
}

int H264RtmpPuser::StartPush()
{
    mIfStart = true;
    this->mThread = new std::thread(ThreadEncode,this);
    mThreadId = this->mThread->get_id();
    return 0;
}

int H264RtmpPuser::StopPush()
{
    mIfConnected = false;
    mIfStart = false;
    if(mThread != nullptr)
        this->mThread->join();
    this->RTMP264_Close();
    return 0;
}