#include "cdyd_ps.h"

#include <stdlib.h>

#include "minivideo/demxer/mpeg/pes/pes_struct.h"
#include "minivideo/demxer/mpeg/pes/pes.h"
#include "minivideo/demxer/mpeg/ps/ps.h"
#include "minivideo/demxer/mpeg/ps/ps_struct.h"
#include "minivideo/bitstream.h"
#include "minivideo/minivideo_mediafile.h"
#include "minivideo/bitstream_utils.h"

#include "minivideo/demxer/mpeg/pes/pes.h"
#include "minivideo/demxer/mpeg/ps/ps.h"

#include "struct.h"

static int parse_system_header(Bitstream_t* bitstr, PesHeader_t* header, SystemHeader_t* packet)
{
    int retcode = SUCCESS;

    TRACE_INFO(MPS, BLD_GREEN "parse_system_header()" CLR_RESET " @ %lli",
        header->offset_start);

    MARKER_BIT
        packet->rate_bound = read_bits(bitstr, 22);
    MARKER_BIT
        packet->audio_bound = read_bits(bitstr, 6);
    packet->fixed_flag = read_bits(bitstr, 1);
    packet->CSPS_flag = read_bits(bitstr, 1);
    packet->system_audio_lock_flag = read_bits(bitstr, 1);
    packet->system_video_lock_flag = read_bits(bitstr, 1);
    MARKER_BIT
        packet->video_bound = read_bits(bitstr, 5);
    packet->packet_rate_restriction_flag = read_bits(bitstr, 1);

    /*unsigned reserved_bits =*/ read_bits(bitstr, 7);

    // stack it?
    while (next_bit(bitstr) == 1)
    {
        packet->stream_id = read_bits(bitstr, 8);
        MARKER_BIT
            MARKER_BIT
            packet->PSTD_buffer_bound_scale = read_bits(bitstr, 1);
        packet->PSTD_buffer_size_bound = read_bits(bitstr, 13);
    }

    return retcode;
}

static int parse_program_stream_directory(Bitstream_t* bitstr, PesHeader_t* header, ProgramStreamDirectory_t* packet)
{
    TRACE_INFO(MPS, BLD_GREEN "parse_program_stream_directory()" CLR_RESET " @ %lli",
        header->offset_start);
    int retcode = SUCCESS;

    packet->number_of_access_units = read_bits(bitstr, 15);
    MARKER_BIT

        packet->prev_directory_offset = read_bits(bitstr, 15) << 30;
    MARKER_BIT
        packet->prev_directory_offset += read_bits(bitstr, 15) << 15;
    MARKER_BIT
        packet->prev_directory_offset += read_bits(bitstr, 15);
    MARKER_BIT

        packet->next_directory_offset = read_bits(bitstr, 15) << 30;
    MARKER_BIT
        packet->next_directory_offset += read_bits(bitstr, 15) << 15;
    MARKER_BIT
        packet->next_directory_offset += read_bits(bitstr, 15);
    MARKER_BIT

        for (uint16_t i = 0; i < packet->number_of_access_units; i++)
        {
            // TODO stack it?
            packet->packet_stream_id = read_bits(bitstr, 8);
            packet->PES_header_position_offset_sign = read_bit(bitstr);

            packet->PES_header_position_offset = read_bits(bitstr, 14) << 30;
            MARKER_BIT
                packet->PES_header_position_offset += read_bits(bitstr, 15) << 15;
            MARKER_BIT
                packet->PES_header_position_offset += read_bits(bitstr, 15);
            MARKER_BIT

                packet->reference_offset = read_bits(bitstr, 16);
            MARKER_BIT
                /*unsigned reserved1 =*/ read_bits(bitstr, 3);

            packet->PTS = read_bits(bitstr, 3) << 30;
            MARKER_BIT
                packet->PTS += read_bits(bitstr, 15) << 15;
            MARKER_BIT
                packet->PTS += read_bits(bitstr, 15);
            MARKER_BIT

                packet->byes_to_read = read_bits(bitstr, 15) << 15;
            MARKER_BIT
                packet->byes_to_read += read_bits(bitstr, 8);
            MARKER_BIT

                packet->intra_coded_indicator = read_bit(bitstr);
            packet->coding_parameters_indicator = read_bits(bitstr, 2);
            /*unsigned reserved2 =*/ read_bits(bitstr, 4);
        }

    return retcode;
}

static int parse_program_stream_map(Bitstream_t* bitstr, PesHeader_t* header, ProgramStreamMap_t* packet)
{
    TRACE_INFO(MPS, BLD_GREEN "parse_program_stream_map()" CLR_RESET " @ %lli",
        header->offset_start);
    int retcode = SUCCESS;
    int i = 0, N1 = 0, N2 = 0;

    packet->current_next_indicator = read_bit(bitstr);
    /*int reserved1 =*/ read_bits(bitstr, 2);
    packet->program_stream_map_version = read_bits(bitstr, 5);
    /*int reserved2 =*/ read_bits(bitstr, 7);
    MARKER_BIT
        packet->program_stream_map_info_length = read_bits(bitstr, 16);
    for (i = 0; i < packet->program_stream_map_info_length; i++)
    {
        // descriptor()
        read_bits(bitstr, 8);
    }

    packet->elementary_stream_map_length = read_bits(bitstr, 16);

    // skip elementary_stream_info
    //read_bits(bitstr, packet->elementary_stream_map_length * 8);

    for (i = 0; i < packet->elementary_stream_map_length; i++)
    {

        // Stack it?
        packet->elementary_stream_map[packet->elementary_stream_map_size].stream_type = read_bits(bitstr, 8);
        i += 1;

        packet->elementary_stream_map[packet->elementary_stream_map_size].elementary_stream_id = read_bits(bitstr, 8);
        i += 1;

        packet->elementary_stream_map[packet->elementary_stream_map_size].elementary_stream_info_length = read_bits(bitstr, 16);
        i += 2;

        read_bits(bitstr, packet->elementary_stream_map[packet->elementary_stream_map_size].elementary_stream_info_length * 8);
        i += packet->elementary_stream_map[packet->elementary_stream_map_size].elementary_stream_info_length;

        //for (i = 0; i < packet->elementary_stream_map[packet->elementary_stream_map_size].elementary_stream_info_length; i++)
        //{
        //    read_bits(bitstr, 8);
        //    // descriptor()
        //}

        packet->elementary_stream_map_size++;

    }

    packet->CRC_32 = read_bits(bitstr, 32);

    return retcode;
}

static int parse_pack_header(Bitstream_t* bitstr, PesHeader_t* header, PackHeader_t* packet)
{
    TRACE_INFO(MPS, BLD_GREEN "parse_pack_header()" CLR_RESET " @ %lli",
        header->offset_start);

    int retcode = SUCCESS;

    // Pack Headers do not have lengh field, rewind 2 bytes
    rewind_bits(bitstr, 16);

    if (read_bits(bitstr, 2) != 1)
    {
        TRACE_ERROR(MPS, "wrong 'marker_bits'");
        printf("offset: %lld\n", bitstream_get_absolute_byte_offset(bitstr));
        return FAILURE;
    }

    packet->system_clock_reference_base = read_bits(bitstr, 3) << 30;
    MARKER_BIT
        packet->system_clock_reference_base += read_bits(bitstr, 15) << 15;
    MARKER_BIT
        packet->system_clock_reference_base += read_bits(bitstr, 15);
    MARKER_BIT
        packet->system_clock_reference_extension = read_bits(bitstr, 9);
    MARKER_BIT
        packet->program_mux_rate = read_bits(bitstr, 22);
    MARKER_BIT
        MARKER_BIT
        /*unsigned reserved =*/ read_bits(bitstr, 5);

    packet->pack_stuffing_length = read_bits(bitstr, 3);

    // Stuffing
    for (uint8_t i = 0; i < packet->pack_stuffing_length; i++)
    {
        // hikvision doesnot know, just skip
        read_bits(bitstr, 8);

        //if (read_bits(bitstr, 8) != 0xFF)
        //{
        //    TRACE_ERROR(MPS, "Wrong 'stuffing_byte'");
        //    return FAILURE;
        //}
    }

    // Pack header have no length field, so we just have to parse them correctly
    header->offset_end = bitstream_get_absolute_byte_offset(bitstr);
    header->payload_length = header->offset_end - header->offset_start - 4;

    return retcode;
}

static int parse_packet(Bitstream_t* bitstr, PesPacket_t* packet) {

    uint32_t skip_length = 0;

    read_bits(bitstr, 2);

    packet->PES_scrambling_control = read_bits(bitstr, 2);
    packet->PES_priority = read_bit(bitstr);
    packet->data_alignment_indicator = read_bit(bitstr);
    packet->copyright = read_bit(bitstr);
    packet->original_or_copy = read_bit(bitstr);
    packet->PTS_DTS_flag = read_bits(bitstr, 2);
    packet->ESCR_flag = read_bit(bitstr);
    packet->ES_rate_flag = read_bit(bitstr);
    packet->DSM_trick_mode_flag = read_bit(bitstr);
    packet->additional_copy_info_flag = read_bit(bitstr);
    packet->PES_CRC_flag = read_bit(bitstr);
    packet->PES_extension_flag = read_bit(bitstr);
    packet->PES_header_data_length = read_bits(bitstr, 8);


    if (packet->PTS_DTS_flag == 2)
    {
        if (read_bits(bitstr, 4) != 2)
        {
            TRACE_ERROR(MPS, "wrong 'marker_bit'");
            return FAILURE;
        }
        packet->PTS = read_bits(bitstr, 3) << 30;
        MARKER_BIT
            packet->PTS += read_bits(bitstr, 15) << 15;
        MARKER_BIT
            packet->PTS += read_bits(bitstr, 15);
        MARKER_BIT

            skip_length += 5;
    }
    else if (packet->PTS_DTS_flag == 3)
    {
        if (read_bits(bitstr, 4) != 3)
        {
            TRACE_ERROR(MPS, "wrong 'marker_bit'");
            return FAILURE;
        }
        packet->PTS = read_bits(bitstr, 3) << 30;
        MARKER_BIT
            packet->PTS += read_bits(bitstr, 15) << 15;
        MARKER_BIT
            packet->PTS += read_bits(bitstr, 15);
        MARKER_BIT
            if (read_bits(bitstr, 4) != 1)
            {
                TRACE_ERROR(MPS, "wrong 'marker_bit'");
                return FAILURE;
            }
        packet->DTS = read_bits(bitstr, 3) << 30;
        MARKER_BIT
            packet->DTS += read_bits(bitstr, 15) << 15;
        MARKER_BIT
            packet->DTS += read_bits(bitstr, 15);
        MARKER_BIT

            skip_length += 10;
    }

    read_bits(bitstr, 8 * (packet->PES_header_data_length - skip_length));

    return 1;
}

static int callback_open_codec(DataCallback* dct, ProgramStreamMap_t& pes_streammap) {

    for (uint8_t i = 0; i < pes_streammap.elementary_stream_map_size; i++) {
        ElementaryStreamMap_t es = pes_streammap.elementary_stream_map[i];
        switch (es.elementary_stream_id)
        {
        case 0xC1: case 0xC2: case 0xC3: case 0xC4: case 0xC5: case 0xC6:
        case 0xC7: case 0xC8: case 0xC9: case 0xCA: case 0xCB: case 0xCC:
        case 0xCD: case 0xCE: case 0xCF: case 0xD0: case 0xD1: case 0xD2:
        case 0xD3: case 0xD4: case 0xD5: case 0xD6: case 0xD7: case 0xD8:
        case 0xD9: case 0xDA: case 0xDB: case 0xDC: case 0xDD: case 0xDE:
        case 0xDF: case SID_AUDIO:
            if (dct) {
                int type = AUDIO_UNKNOWN;
                switch (es.stream_type)
                {
                case 0x90: type = AUDIO_G711_A; break;
                case 0x91: type = AUDIO_G711_MU; break;
                case 0x92: type = AUDIO_G722; break;
                case 0x93: type = AUDIO_G723_1; break;
                case 0x99: type = AUDIO_G729; break;
                case 0x9B: break;
                default:
                    break;
                }
                dct->open_codec(type, AV_AUDIO);
            }
            break;
        case 0xE1: case 0xE2: case 0xE3: case 0xE4: case 0xE5:
        case 0xE6: case 0xE7: case 0xE8: case 0xE9: case 0xEA:
        case 0xEB: case 0xEC: case 0xED: case 0xEE: case 0xEF:
        case SID_VIDEO:
            if (dct) {
                int type = VIDEO_UNKNOWN;
                switch (es.stream_type)
                {
                case 0x1B: type = VIDEO_H264; break;
                case 0x10: type = VIDEO_MPEG4; break;
                case 0x80: break;
                default:
                    break;
                }
                dct->open_codec(type, AV_VIDEO);
            }
            break;
        default:
            break;
        }
    }
    return 1;
}

int cdyd_PS_fileParse(MediaFile_t* media, DataCallback* dct) {
    int retcode = SUCCESS;

    TRACE_INFO(MPS, BLD_GREEN "ps_fileParse()" CLR_RESET);

    // Init bitstream to parse container infos
    Bitstream_t* bitstr = init_bitstream(media);

    if (bitstr != NULL)
    {
        // Init an MpegPs structure
        MpegPs_t mpg;
        memset(&mpg, 0, sizeof(MpegPs_t));

        // A convenient way to stop the parser
        mpg.run = true;

        // stuff
        const int64_t min_packet_size = 4;

        int64_t av_length = 0;

        uint32_t start_code = 0;

        // Loop on PES packets
        while (mpg.run == true &&
            retcode == SUCCESS &&
            bitstream_get_absolute_byte_offset(bitstr) < (media->file_size - min_packet_size))
        {
            if (dct && dct->is_quit()) break;

            start_code = next_bits(bitstr, 24);
            if (start_code != 0x000001) {
                skip_bits(bitstr, 8);
                continue;
            }
            int64_t abs_offset = bitstream_get_absolute_byte_offset(bitstr);

            // Init
            PackHeader_t pack_header;
            memset(&pack_header, 0, sizeof(PackHeader_t));
            SystemHeader_t system_header;
            memset(&system_header, 0, sizeof(SystemHeader_t));

            PesHeader_t pes_header;
            memset(&pes_header, 0, sizeof(PesHeader_t));
            PesPacket_t pes_packet;
            memset(&pes_packet, 0, sizeof(PesPacket_t));
            ProgramStreamMap_t pes_streammap;
            memset(&pes_streammap, 0, sizeof(ProgramStreamMap_t));
            ProgramStreamDirectory_t pes_streamdirectory;
            memset(&pes_streamdirectory, 0, sizeof(ProgramStreamDirectory_t));

            // Parse packet header
            parse_pes_header(bitstr, &pes_header);

            switch (pes_header.stream_id)
            {
            case SID_PACK_HEADER:
                retcode = parse_pack_header(bitstr, &pes_header, &pack_header);
                mpg.stat_packheader++;
                break;
            case SID_SYSTEM_HEADER:
                retcode = parse_system_header(bitstr, &pes_header, &system_header);
                mpg.stat_systemheader++;
                break;

            case SID_PROGRAM_STREAM_MAP:
            {
                retcode = parse_program_stream_map(bitstr, &pes_header, &pes_streammap);

                if (retcode) {
                    callback_open_codec(dct, pes_streammap);
                }

                mpg.stat_packet_psm++;

            }
            break;
            case SID_PROGRAM_STREAM_DIRECTORY:
                retcode = parse_program_stream_directory(bitstr, &pes_header, &pes_streamdirectory);
                mpg.stat_packet_psd++;
                break;
            case SID_PRIVATE_STREAM_2:
                TRACE_2(MPS, BLD_GREEN "Private Stream 2 PES" CLR_RESET " @ %lli", pes_header.offset_start);
                mpg.stat_packet_private++;
                break;
            case SID_PADDING:
                retcode = parse_pes_padding(bitstr, &pes_header, &pes_packet);
                mpg.stat_packet_other++;
                break;
            case SID_PRIVATE_STREAM_1:
                TRACE_2(MPS, BLD_GREEN "Private Stream 1 PES" CLR_RESET " @ %lli", pes_header.offset_start);
                mpg.stat_packet_private++;
                break;
            case 0xC1: case 0xC2: case 0xC3: case 0xC4: case 0xC5: case 0xC6:
            case 0xC7: case 0xC8: case 0xC9: case 0xCA: case 0xCB: case 0xCC:
            case 0xCD: case 0xCE: case 0xCF: case 0xD0: case 0xD1: case 0xD2:
            case 0xD3: case 0xD4: case 0xD5: case 0xD6: case 0xD7: case 0xD8:
            case 0xD9: case 0xDA: case 0xDB: case 0xDC: case 0xDD: case 0xDE:
            case 0xDF: case SID_AUDIO:
            {
                TRACE_INFO(MPS, BLD_GREEN "parse_pes_audio()" CLR_RESET " @ %lli", pes_header.offset_start);

                parse_packet(bitstr, &pes_packet);

                bitstream_force_alignment(bitstr);

                uint32_t skip_length = pes_header.offset_end - bitstream_get_absolute_byte_offset(bitstr);
                if (skip_length > media->file_size) {
                    retcode = FAILURE;
                    break;
                }
                av_length += skip_length;

                uint8_t* chunk = (uint8_t*)malloc(skip_length);
                read_chunk(bitstr, chunk, skip_length);
                if (dct) {
                    packet_t* pkt = (packet_t*)malloc(sizeof(packet_t));
                    pkt->data = chunk;
                    pkt->length = skip_length;
                    pkt->pts = pes_packet.PTS;
                    pkt->dts = pes_packet.DTS;
                    pkt->type = AV_AUDIO;
                    dct->push_packet(pkt, AV_AUDIO);
                }

            }
            break;

            case 0xE1: case 0xE2: case 0xE3: case 0xE4: case 0xE5:
            case 0xE6: case 0xE7: case 0xE8: case 0xE9: case 0xEA:
            case 0xEB: case 0xEC: case 0xED: case 0xEE: case 0xEF:
            case SID_VIDEO:
            {
                TRACE_INFO(MPS, BLD_GREEN "parse_pes_video()" CLR_RESET " @ %lli", pes_header.offset_start + 6);

                parse_packet(bitstr, &pes_packet);

                bitstream_force_alignment(bitstr);

                uint32_t skip_length = pes_header.offset_end - bitstream_get_absolute_byte_offset(bitstr);
                if (skip_length > media->file_size) {
                    retcode = FAILURE;
                    break;
                }
                av_length += skip_length;

                uint8_t* chunk = (uint8_t*)malloc(skip_length);
                read_chunk(bitstr, chunk, skip_length);

                if (dct) {
                    packet_t* pkt = (packet_t*)malloc(sizeof(packet_t));
                    pkt->data = chunk;
                    pkt->length = skip_length;
                    pkt->pts = pes_packet.PTS;
                    pkt->dts = pes_packet.DTS;
                    pkt->type = AV_VIDEO;
                    dct->push_packet(pkt, AV_VIDEO);
                }

            }
            break;

            case SID_PROGRAM_END:
                mpg.stat_packet_other++;
                mpg.run = false;
                break;

            default:
                TRACE_WARNING(MPS, "Unknown PES packet type (0x%02X) @ %lli",
                    pes_header.stream_id, pes_header.offset_start);
                mpg.stat_packet_other++;
                // find next startcode
                uint32_t next_startcode = -1;
                while (bitstream_get_absolute_byte_offset(bitstr) < (media->file_size - min_packet_size)) {
                    next_startcode = next_bits(bitstr, 24);
                    if (next_startcode == 0x000001) {
                        break;
                    }
                    skip_bits(bitstr, 8);
                }
                continue;
            }

            retcode = jumpy_pes(bitstr, &pes_header);
        }

        // Free bitstream
        free_bitstream(&bitstr);
        if (dct) dct->push_finish();
        printf("demuxer over!!!");
        if (av_length < media->file_size / 2) {
            retcode = FAILURE;
        }

        // Recap
        TRACE_INFO(MPS, "MPEG PS (version %u) stats", mpg.mpeg_version);
        TRACE_INFO(MPS, "- Pack Headers:    %u", mpg.stat_packheader);
        TRACE_INFO(MPS, "- System Headers:  %u", mpg.stat_systemheader);
        TRACE_INFO(MPS, "- PSM packets:     %u", mpg.stat_packet_psm);
        TRACE_INFO(MPS, "- PSD packets:     %u", mpg.stat_packet_psd);
        TRACE_INFO(MPS, "- Private packets: %u", mpg.stat_packet_private);
        TRACE_INFO(MPS, "- Audio packets:   %u", mpg.stat_packet_audio);
        TRACE_INFO(MPS, "- Video packets:   %u", mpg.stat_packet_video);
        TRACE_INFO(MPS, "- Unknown packets: %u", mpg.stat_packet_other);
    }
    else
    {
        retcode = FAILURE;
    }

    return retcode;
}