#include "h2645_parser.h"

#include <stdlib.h>
#include <stdio.h>

#include "h264_stream.h"
#include "h265_stream.h"

#include "struct.h"

const int MAX_NAL_SIZE = 1 * 1024 * 1024;

inline int findStartcode3(unsigned char* buffer)
{
    return (buffer[0] == 0 && buffer[1] == 0 && buffer[2] == 1);
}

inline int findStartcode4(unsigned char* buffer)
{
    return (buffer[0] == 0 && buffer[1] == 0 && buffer[2] == 0 && buffer[3] == 1);
}

static int find_first_NALU(FILE* fp, int *start_code_length) {
    int found = 0;
    int info2 = 0;
    int info3 = 0;
    int eof = 0;
    int pos = 0;
    int startcode_len = 0;
    unsigned char* buffer = NULL;

    if ((buffer = (unsigned char*)calloc(MAX_NAL_SIZE, sizeof(char))) == NULL)
        printf("Could not allocate buffer memory\n");

    while (!found && !feof(fp))
    {
        buffer[pos++] = fgetc(fp);//¶ÁÒ»¸ö×Ö½Úµ½BUFÖÐ

        info3 = findStartcode4(&buffer[pos - 4]);//ÅжÏÊÇ·ñΪ0x00000001
        if (info3 != 1)
        {
            info2 = findStartcode3(&buffer[pos - 3]);//ÅжÏÊÇ·ñΪ0x000001
            if (info2)
            {
                startcode_len = 3;
            }
        }
        else
        {
            startcode_len = 4;
        }

        found = (info2 == 1 || info3 == 1);
        if (pos >= MAX_NAL_SIZE)
        {
            free(buffer);
            return -1;
        }
    }

    // ÎļþÖ¸ÕëÒª»Ö¸´
    fseek(fp, -startcode_len, SEEK_CUR);

    free(buffer);
    if (start_code_length != NULL)
        *start_code_length = startcode_len;

    return pos - startcode_len;
}

H2645 probe_h2645(FILE* fp) {
    int offset = 0;
    int startcode = 0;
    unsigned char nalHader = 0;
    unsigned char nalType = 0;

    H2645 type = TYPE_H264; // default

    offset = find_first_NALU(fp, &startcode);
    if (offset < 0)
    {
        return TYPE_UNKNOWN;
    }
    fseek(fp, offset + startcode, SEEK_SET);
    fread((void*)&nalHader, 1, 1, fp);
    // check h264 first...
    nalType = nalHader & 0x1f; // 5 bit
    if (nalType > 0 && nalType < 22) // ok
    {
        type = TYPE_H264;
    }
    else
    {
        // not h264, then check h265...
        nalType = (nalHader >> 1) & 0x3f; // 6 bit
        if (nalType >= 0 && nalType <= 47) // ok
        {
            type = TYPE_H265;
        }
    }
    fseek(fp, 0, SEEK_SET);

    return type;
}

typedef struct
{
    int type;                       // 0 -- h.264; 1 -- h.265
    unsigned int num;               // ÐòºÅ
    unsigned int len;               // º¬ÆðʼÂëµÄ×ܵij¤¶È
    unsigned int offset;       // nal°üÔÚÎļþÖеÄÆ«ÒÆ
    int sliceType;               // Ö¡ÀàÐÍ
    int nalType;            // NALÀàÐÍ
    int startcodeLen;             // start code³¤¶È
    char startcodeBuffer[16];         // ÆðʼÂ룬×Ö·û´®ÐÎʽ
} NALU_t;

typedef struct
{
    int profile_idc;
    int level_idc;
    int width;
    int height;
    int crop_left;
    int crop_right;
    int crop_top;
    int crop_bottom;
    float max_framerate;  // ÓÉSPS¼ÆËãµÃµ½µÄÖ¡ÂÊ£¬Îª0±íʾSPSÖÐûÓÐÏàÓ¦µÄ×ֶμÆËã
    int chroma_format_idc;  // YUVÑÕÉ«¿Õ¼ä 0: monochrome 1:420 2:422 3:444
}SPSInfo_t;

typedef struct
{
    int encoding_type;  // Ϊ1±íʾCABAC 0±íʾCAVLC

}PPSInfo_t;

static int getAnnexbNALU(FILE* fp, NALU_t* nalu, 
    h264_stream_t *h264, h265_stream_t *hevc,
    uint8_t **wholeNal)
{
    int pos = 0;
    int found, rewind;
    unsigned char* buffer;
    int info2 = 0, info3 = 0;
    int eof = 0;

    if ((buffer = (unsigned char*)calloc(MAX_NAL_SIZE, sizeof(char))) == NULL)
        printf("Could not allocate buffer memory\n");

    if (3 != fread(buffer, 1, 3, fp))//´ÓÂëÁ÷ÖжÁ3¸ö×Ö½Ú
    {
        free(buffer);
        return 0;
    }
    info2 = findStartcode3(buffer);//ÅжÏÊÇ·ñΪ0x000001
    if (info2 != 1)
    {
        //Èç¹û²»ÊÇ£¬ÔÙ¶ÁÒ»¸ö×Ö½Ú
        if (1 != fread(buffer + 3, 1, 1, fp))//¶ÁÒ»¸ö×Ö½Ú
        {
            free(buffer);
            return 0;
        }
        info3 = findStartcode4(buffer);//ÅжÏÊÇ·ñΪ0x00000001
        if (info3 != 1)//Èç¹û²»ÊÇ£¬·µ»Ø-1
        {
            free(buffer);
            return -1;
        }
        else
        {
            //Èç¹ûÊÇ0x00000001,µÃµ½¿ªÊ¼Ç°×ºÎª4¸ö×Ö½Ú
            nalu->startcodeLen = 4;
        }
    }
    else
    {
        //Èç¹ûÊÇ0x000001,µÃµ½¿ªÊ¼Ç°×ºÎª3¸ö×Ö½Ú
        nalu->startcodeLen = 3;
    }

    pos = nalu->startcodeLen;
    //²éÕÒÏÂÒ»¸ö¿ªÊ¼×Ö·ûµÄ±ê־λ
    found = 0;
    info2 = 0;
    info3 = 0;

    while (!found)
    {
        if (feof(fp))//ÅжÏÊÇ·ñµ½ÁËÎļþβ
        {
            eof = 1;
            goto got_nal;
        }
        buffer[pos++] = fgetc(fp);//¶ÁÒ»¸ö×Ö½Úµ½BUFÖÐ

        info3 = findStartcode4(&buffer[pos - 4]);//ÅжÏÊÇ·ñΪ0x00000001
        if (info3 != 1)
            info2 = findStartcode3(&buffer[pos - 3]);//ÅжÏÊÇ·ñΪ0x000001

        found = (info2 == 1 || info3 == 1);
    }

    // startcode¿ÉÄÜΪ3£¬Ò²¿ÉÄÜΪ4£¬¹ÊÒªÈç´ËÅжÏ
    rewind = (info3 == 1) ? -4 : -3;

    if (0 != fseek(fp, rewind, SEEK_CUR))//°ÑÎļþÖ¸ÕëÖ¸ÏòÇ°Ò»¸öNALUµÄĩβ
    {
        free(buffer);
        printf("Cannot fseek in the bit stream file");
    }

got_nal:
    // µ±´ïµ½Îļþĩβʱ£¬»ØÍË1¸öλÖÃ
    if (eof)
    {
        rewind = -1;
    }

    // °üÀ¨ÆðʼÂëÔÚÄÚµÄ5¸ö×Ö½Ú
    if (nalu->startcodeLen == 3)
        sprintf(nalu->startcodeBuffer, "%02x%02x%02x%02x", buffer[0], buffer[1], buffer[2], buffer[3]);
    else
        sprintf(nalu->startcodeBuffer, "%02x%02x%02x%02x%02x", buffer[0], buffer[1], buffer[2], buffer[3], buffer[4]);
    nalu->len = pos + rewind;

    uint8_t nal_header = 0;
    if (nalu->type)
    {
        hevc->sh->read_slice_type = 1;
        h265_read_nal_unit(hevc, &buffer[nalu->startcodeLen], nalu->len - nalu->startcodeLen);
        nalu->nalType = hevc->nal->nal_unit_type;
        nalu->sliceType = hevc->sh->slice_type;
        hevc->sh->read_slice_type = 0;
    }
    else
    {
        // simple version
#if 0
        nal_header = buffer[nalu->startcodeLen];
        nalu->nalType = nal_header & 0x1f;// 5 bit

        // »ñÈ¡sliceÀàÐÍ£ºIÖ¡¡¢PÖ¡¡¢BÖ¡
        // ×¢£ºÔÚnalÀàÐÍΪ1~5ʱ»ñÈ¡
        if (nalu->nalType <= 5 && nalu->nalType >= 1)
        {
            int start_bit = 0;
            int first_mb_in_slice = ue((char*)buffer + nalu->startcodeLen + 1, 8, start_bit);
            nalu->sliceType = ue((char*)buffer + nalu->startcodeLen + 1, 8, start_bit);
        }
        if (nalu->nalType == 7 || nalu->nalType == 8) // sps pps
        {
            read_nal_unit(m_hH264, &buffer[nalu->startcodeLen], nalu->len - nalu->startcodeLen);
        }
#else
        uint8_t* nal_buf = (uint8_t*)malloc(nalu->len);
        memcpy(nal_buf, buffer, nalu->len);
        *wholeNal = nal_buf;

        h264->sh->read_slice_type = 1;
        read_nal_unit(h264, &buffer[nalu->startcodeLen], nalu->len - nalu->startcodeLen);
        nalu->nalType = h264->nal->nal_unit_type;
        nalu->sliceType = h264->sh->slice_type;
        h264->sh->read_slice_type = 0;
#endif
    }

    free(buffer);

    return (pos + rewind);//·µ»ØÁ½¸ö¿ªÊ¼×Ö·ûÖ®¼ä¼ä¸ôµÄ×Ö½ÚÊý£¬¼´°üº¬ÓÐǰ׺µÄNALUµÄ³¤¶È
}

static void showNalInfo(NALU_t* nalu);

static int nalu_type(NALU_t* nalu) {
    if (nalu->type == 0)
    {
        // NALµ¥ÔªÀàÐÍ
        switch (nalu->nalType)
        {
        case 0: case 1: case 2: case 3: case 4:case 5:case 6:
        case 7:case 8:case 9:case 10:case 11:case 12:case 13:
        case 19:
            return nalu->nalType;
        default:
            return -1;
        }
    }
    else
    {
        // NALµ¥ÔªÀàÐÍ
        switch (nalu->nalType)
        {
            // to confirm type...
        case NAL_UNIT_CODED_SLICE_TRAIL_N:
        case NAL_UNIT_CODED_SLICE_TRAIL_R:
        case NAL_UNIT_CODED_SLICE_TSA_N:
        case NAL_UNIT_CODED_SLICE_TSA_R:
        case NAL_UNIT_CODED_SLICE_RADL_N:
        case NAL_UNIT_CODED_SLICE_RADL_R:
        case NAL_UNIT_CODED_SLICE_IDR_W_RADL:
        case NAL_UNIT_CODED_SLICE_IDR_N_LP:
        case NAL_UNIT_CODED_SLICE_CRA:
        case NAL_UNIT_PREFIX_SEI:
        case NAL_UNIT_SUFFIX_SEI:
        case NAL_UNIT_VPS:
        case NAL_UNIT_SPS:
        case NAL_UNIT_PPS:
        case NAL_UNIT_AUD:
        case NAL_UNIT_EOS:
        case NAL_UNIT_EOB:
        case NAL_UNIT_FILLER_DATA:
            return nalu->nalType;
        default:
            return -1;
        }
    }
    return -1;
}

int h2645_fileParse(FILE* fp, DataCallback* dct) {

    H2645 type = probe_h2645(fp);
    if (type == TYPE_UNKNOWN) {
        return 0;
    }

    h264_stream_t* h264_s = NULL;
    h265_stream_t* hevc_s = NULL;

    if (type == TYPE_H264) {
        h264_s = h264_new();
    }
    if (type == TYPE_H265) {
        hevc_s = h265_new();
    }

    NALU_t n;
    int nal_num = 0;
    int offset = 0;
    int nalLen = 0;

    memset(&n, '\0', sizeof(NALU_t));

    int startcodelen = 0;
    offset = find_first_NALU(fp, &startcodelen);

    if (offset < 0)
    {
        return 0;
    }

    fseek(fp, offset, SEEK_SET);

    uint8_t* nalu_buf = NULL;


    while (!feof(fp))
    {
        if (dct && dct->is_quit()) break;

        //printf("\n\n");
        nalLen = getAnnexbNALU(fp, &n, h264_s, hevc_s, &nalu_buf);//ÿִÐÐÒ»´Î£¬ÎļþµÄÖ¸ÕëÖ¸Ïò±¾´ÎÕÒµ½µÄNALUµÄĩ⣬ÏÂÒ»¸öλÖü´ÎªÏ¸öNALUµÄÆðʼÂë0x000001
        n.offset = offset;
        n.num = nal_num;
        offset = offset + nalLen;

        //showNalInfo(&n);
        if (n.len > 0) {
            if (dct) {
                
                packet_t* pkt = (packet_t*)malloc(sizeof(packet_t));
                pkt->data = nalu_buf;
                pkt->length = n.len;
                pkt->pts = -1;
                pkt->dts = -1;

                if (nalu_type(&n) >= 0) {
                    pkt->type = AV_VIDEO;
                }
                else
                {
                    pkt->type = AV_AUDIO;
                }
                dct->push_packet(pkt, pkt->type);
                
            }
        }
        //printf("\n\n");

        nal_num++;
    }

    return 0;
}

static void showNalInfo(NALU_t* nalu) {
    if (nalu->type == 0)
    {
        // NALµ¥ÔªÀàÐÍ
        switch (nalu->nalType)
        {
        case 0:
            printf("Unspecified\n");
            break;
        case 1:
            printf("Coded slice of a non-IDR picture\n");
            switch (nalu->sliceType)
            {
            case 0:
            case 5:
                printf("P Slice #%d\n");
                break;
            case 1:
            case 6:
                printf("B Slice #%d\n");
                break;
            case 2:
            case 7:
                printf("I Slice #%d\n");
                break;
            }
            break;
        case 2:
            printf("DPA\n");
            break;
        case 3:
            printf("DPB\n");
            break;
        case 4:
            printf("DPC\n");
            break;
        case 5:
            printf("Coded slice of an IDR picture\n");
            printf("IDR #%d\n");
            break;
        case 6:
            printf("Supplemental enhancement information\n");
            printf("SEI\n");
            break;
        case 7:
            printf("Sequence parameter set\n");
            printf("SPS\n");
            break;
        case 8:
            printf("Picture parameter set\n");
            printf("PPS\n");
            break;
        case 9:
            printf("Access UD\n");
            printf("AUD\n");
            break;
        case 10:
            printf("END_SEQUENCE\n");
            break;
        case 11:
            printf("END_STREAM\n");
            break;
        case 12:
            printf("FILLER_DATA\n");
            break;
        case 13:
            printf("SPS_EXT\n");
            break;
        case 19:
            printf("AUXILIARY_SLICE\n");
            break;
        default:
            printf("Other\n");
            break;
        }
    }
    else
    {
        // NALµ¥ÔªÀàÐÍ
        switch (nalu->nalType)
        {
            // to confirm type...
        case NAL_UNIT_CODED_SLICE_TRAIL_N:
        case NAL_UNIT_CODED_SLICE_TRAIL_R:
            printf("Coded slice segment of a non-TSA, non-STSA trailing picture\n");
            switch (nalu->sliceType)
            {
            case H265_SH_SLICE_TYPE_B:
                printf("B Slice #%d\n");
                break;
            case H265_SH_SLICE_TYPE_P:
                printf("P Slice #%d\n");
                break;
            case H265_SH_SLICE_TYPE_I:
                printf("I Slice #%d\n");
                break;
            }
            break;
        case NAL_UNIT_CODED_SLICE_TSA_N:
        case NAL_UNIT_CODED_SLICE_TSA_R:
            printf("Coded slice segment of a TSA picture\n");
            switch (nalu->sliceType)
            {
            case H265_SH_SLICE_TYPE_B:
                printf("B Slice #%d\n");
                break;
            case H265_SH_SLICE_TYPE_P:
                printf("P Slice #%d\n");
                break;
            case H265_SH_SLICE_TYPE_I:
                printf("I Slice #%d\n");
                break;
            }
            break;
        case NAL_UNIT_CODED_SLICE_RADL_N:
        case NAL_UNIT_CODED_SLICE_RADL_R:
            printf("Coded slice segment of a TSA picture\n");
            switch (nalu->sliceType)
            {
            case H265_SH_SLICE_TYPE_B:
                printf("B Slice #%d\n");
                break;
            case H265_SH_SLICE_TYPE_P:
                printf("P Slice #%d\n");
                break;
            case H265_SH_SLICE_TYPE_I:
                printf("I Slice #%d\n");
                break;
            }
            break;
        case NAL_UNIT_CODED_SLICE_IDR_W_RADL:
        case NAL_UNIT_CODED_SLICE_IDR_N_LP:
            printf("Coded slice of an IDR picture\n");
            printf("IDR #%d\n");
            break;
        case NAL_UNIT_CODED_SLICE_CRA:
            printf("Coded slice segment of a CRA picture\n");
            printf("CRA #%d\n");
            break;
        case NAL_UNIT_PREFIX_SEI:
        case NAL_UNIT_SUFFIX_SEI:
            printf("Supplemental enhancement information\n");
            printf("SEI\n");
            break;
        case NAL_UNIT_VPS:
            printf("Video parameter set\n");
            printf("VPS\n");
            break;
        case NAL_UNIT_SPS:
            printf("Sequence parameter set\n");
            printf("SPS\n");
            break;
        case NAL_UNIT_PPS:
            printf("Picture parameter set\n");
            printf("PPS\n");
            break;
        case NAL_UNIT_AUD:
            printf("Access UD\n");
            printf("AUD\n");
            break;
        case NAL_UNIT_EOS:
            printf("END_SEQUENCE\n");
            break;
        case NAL_UNIT_EOB:
            printf("END_STREAM\n");
            break;
        case NAL_UNIT_FILLER_DATA:
            printf("FILLER_DATA\n");
            break;
        default:
            printf("Unknown\n");
            break;
        }
    }

    // ÐòºÅ
    printf("number: %d\n", nalu->num + 1);  // ÏòÁ¿ÖеÄÐòºÅ´Ó0¿ªÊ¼
    // Êý¾ÝÆ«ÒÆ
    printf("offset: %08X\n", nalu->offset);
    // ³¤¶È
    printf("length: %d\n", nalu->len);
    // ÆðʼÂë
    printf("start code: %s\n", nalu->startcodeBuffer);
}