#include "FlvParser.h"

#include <stdlib.h>

#include <fstream>
#include <iostream>

int FlvParser::AudioTag::_aacProfile;
int FlvParser::AudioTag::_sampleRateIndex;
int FlvParser::AudioTag::_channelConfig;

static int nH264StartCode = 0x01000000;
#define CheckBuf(x)                  \
    {                                \
        if (nBufLen - nOffset < x) { \
            nUsedLen = nOffset;      \
            return 0;                \
        }                            \
    }

FlvParser::FlvParser() { _pFlvHeader = nullptr; }

FlvParser::~FlvParser() {
    for (int i = 0; i < _vpTag.size(); i++) {
        DestroyTag(_vpTag[i]);
        delete _vpTag[i];
    }
}

int FlvParser::Parse(uint8_t* pBuf, int nBufLen, int& nUsedLen) {
    int nOffset = 0;

    if (_pFlvHeader == nullptr) {
        CheckBuf(9);
        _pFlvHeader = CreateHeader(pBuf + nOffset);
        nOffset += _pFlvHeader->nHeaderSize;
    }

    while (1) {
        CheckBuf(15);

        int nPrevSize = SolveU32(pBuf + nOffset);
        nOffset += 4;
        Tag* tag = CreateTag(pBuf + nOffset, nBufLen - nOffset);

        if (tag == nullptr) {
            nOffset -= 4;
            break;
        }

        nOffset += (11 + tag->_header.nDataSize);
        _vpTag.push_back(tag);
    }

    nUsedLen = nOffset;

    return 0;
}

FlvParser::FlvHeader* FlvParser::CreateHeader(uint8_t* pBuf) {
    FlvHeader* pHeader = new FlvHeader;

    pHeader->nVersion = pBuf[3];
    pHeader->bHaveAudio = (pBuf[4] >> 2) & 0x01;
    pHeader->bHaveVideo = (pBuf[4] >> 0) & 0x01;
    pHeader->nHeaderSize = SolveU32(pBuf + 5);
    pHeader->pFlvHeader = new uint8_t[pHeader->nHeaderSize];

    memset(pHeader->pFlvHeader, 0, pHeader->nHeaderSize);
    memcpy(pHeader->pFlvHeader, pBuf, pHeader->nHeaderSize);

    return pHeader;
}

void FlvParser::DestroyHeader(FlvParser::FlvHeader* pHeader) {
    if (pHeader) {
        if (pHeader->pFlvHeader) {
            delete pHeader->pFlvHeader;
        }
    }
}

void FlvParser::Tag::Init(TagHeader* pHeader, uint8_t* pBuf, int nLeftLen) {
    memcpy(&_header, pHeader, sizeof(TagHeader));
    _pTagHeader = new uint8_t[11];
    memcpy(_pTagHeader, pBuf, 11);

    _pTagData = new uint8_t[pHeader->nDataSize];
    memcpy(_pTagData, pBuf + 11, pHeader->nDataSize);
}

FlvParser::VideoTag::VideoTag(TagHeader* pHeader, uint8_t* pBuf, int nLeftLen,
                              FlvParser* pParser) {
    Init(pHeader, pBuf, nLeftLen);

    uint8_t* pd = _pTagData;
    _nFrameType = (pd[0] & 0xf0) >> 4;
    _nCodecType = (pd[0] & 0x0f) >> 0;

    if (_header.nTagType == 0x09 && _nCodecType == 7) {
        ParseH264Tag(pParser);
    }
}

FlvParser::AudioTag::AudioTag(TagHeader* pHeader, uint8_t* pBuf, int nLeftLen,
                              FlvParser* pParser) {
    Init(pHeader, pBuf, nLeftLen);

    uint8_t* pd = _pTagData;
    _nSoundFormat = (pd[0] & 0xf0) >> 4;
    _nSoundRate = (pd[0] & 0x0c) >> 2;
    _nSoundSize = (pd[0] & 0x02) >> 1;
    _nSoundType = (pd[0] & 0x01) >> 0;

    if (_nSoundFormat == 10) {
        ParseAACTag(pParser);
    }
}

int FlvParser::AudioTag::ParseAACTag(FlvParser* pParser) {
    uint8_t* pd = _pTagData;

    int nAACPacketType = pd[1];

    if (nAACPacketType == 0) {
        ParseAudioSpecificConfig(pParser, pd);
    } else if (nAACPacketType == 1) {
        ParseAACRaw(pParser, pd);
    } else {
    }

    return 0;
}

int FlvParser::AudioTag::ParseAudioSpecificConfig(FlvParser* pParser,
                                                  uint8_t* pTagData) {
    uint8_t* pd = pTagData;

    _aacProfile = (pd[2] & 0xf8) >> 3;
    _sampleRateIndex = ((pd[2] & 0x07) << 1 | (pd[3] >> 7));
    _channelConfig = (pd[3] >> 3) & 0x0f;

    printf("-----AAC-----\n");
    printf("aacProfile:%d\n", _aacProfile);
    printf("sampleRateIndex:%d\n", _sampleRateIndex);
    printf("channelConfig:%d\n", _channelConfig);

    _pMedia = nullptr;
    _nMediaLen = 0;

    return 0;
}

int FlvParser::AudioTag::ParseAACRaw(FlvParser* pParser, uint8_t* pTagData) {
    uint64_t bits = 0;
    int datasize = _header.nDataSize - 2;

    WriteU64(bits, 12, 0xFFF);
    WriteU64(bits, 1, 0);
    WriteU64(bits, 2, 0);
    WriteU64(bits, 1, 1);
    WriteU64(bits, 2, _aacProfile - 1);
    WriteU64(bits, 4, _sampleRateIndex);
    WriteU64(bits, 1, 0);
    WriteU64(bits, 3, _channelConfig);
    WriteU64(bits, 1, 0);
    WriteU64(bits, 1, 0);
    WriteU64(bits, 1, 0);
    WriteU64(bits, 1, 0);
    WriteU64(bits, 13, datasize + 7);
    WriteU64(bits, 11, 0x7FF);
    WriteU64(bits, 2, 0);

    _nMediaLen = datasize + 7;
    _pMedia = new uint8_t[_nMediaLen];

    uint8_t p64[8] = {0};

    for (size_t i = 0; i < 8; i++) {
        p64[i] = bits >> (64 - 8 * (i + 1));
    }

    memcpy(_pMedia, p64 + 1, 7);
    memcpy(_pMedia + 7, pTagData + 2, datasize);
}

FlvParser::Tag* FlvParser::CreateTag(uint8_t* pBuf, int nLeftLen) {
    TagHeader* header;
    header->nTagType = SolveU8(pBuf);
    header->nDataSize = SolveU24(pBuf + 1);
    header->nTimeSamp = SolveU24(pBuf + 4);
    header->nTimeSampExt = SolveU8(pBuf + 7);
    header->nStreamID = SolveU24(pBuf + 8);
    header->nTotalTS =
        (uint32_t)(header->nTimeSampExt << 24) + header->nTimeSamp;

    if ((header->nDataSize + 11) < nLeftLen) {
        return nullptr;
    }

    Tag* pTag;

    switch (header->nTagType) {
        case 0x08:
            // video
            pTag = new VideoTag(header, pBuf, nLeftLen, this);
            break;
        case 0x09:
            // audio
            pTag = new AudioTag(header, pBuf, nLeftLen, this);
            break;
        case 0x12:
            // script
        default:
            pTag = new Tag();
            pTag->Init(header, pBuf, nLeftLen);
            break;
    }

    return pTag;
}

void FlvParser::DestroyTag(Tag* pTag) {
    if (pTag->_pTagHeader) delete[] pTag->_pTagHeader;
    if (pTag->_pMedia) delete[] pTag->_pMedia;
    if (pTag->_pTagData) delete[] pTag->_pTagData;
}

int FlvParser::VideoTag::ParseH264Tag(FlvParser* pParser) {
    uint8_t* pd = _pTagData;

    int nAVCPacketType = pd[1];

    int nCompositionTime = SolveU24(pd + 2);

    if (nAVCPacketType == 0) {
        ParseH264Configuration(pParser, pd);
    } else if (nAVCPacketType == 1) {
        ParseNalu(pParser, pd);
    }

    return 0;
}

int FlvParser::VideoTag::ParseH264Configuration(FlvParser* pParser,
                                                uint8_t* pTagData) {
    uint8_t* pd = pTagData;

    pParser->_nNaluUnitLength = (pd[9] & 0x03) + 1;

    int pps_size, sps_size;

    sps_size = SolveU16(pd + 11);

    pps_size = SolveU16(pd + 11 + 2 + (1 + sps_size) + 1);

    _nMediaLen = 4 + sps_size + 4 + pps_size;
    _pMedia = new uint8_t[_nMediaLen];

    memset(_pMedia, 0, _nMediaLen);

    memcpy(_pMedia, &nH264StartCode, 4);
    memcpy(_pMedia + 4, pd + 11 + 2, sps_size);
    memcpy(_pMedia + 4 + sps_size, &nH264StartCode, 4);
    memcpy(_pMedia + 8 + sps_size, pd + 11 + 2 + sps_size + 1 + 2, pps_size);

    return 0;
}

int FlvParser::VideoTag::ParseNalu(FlvParser* pParser, uint8_t* pTagData) {
    uint8_t* pd = pTagData;

    int nOffset = 0;

    _nMediaLen = 0;
    _pMedia = new uint8_t[_header.nDataSize + 10];

    nOffset = 5;

    while (1) {
        if (nOffset >= _header.nDataSize) break;

        int nNaluLen = 0;

        switch (pParser->_nNaluUnitLength) {
            case 4:
                nNaluLen = SolveU32(pd + nOffset);
                break;
            case 3:
                nNaluLen = SolveU24(pd + nOffset);
                break;
            case 2:
                nNaluLen = SolveU16(pd + nOffset);
                break;
            default:
                nNaluLen = SolveU8(pd + nOffset);
                break;
        }

        memcpy(_pMedia + _nMediaLen, &nH264StartCode, 4);
        memcpy(_pMedia + _nMediaLen + 4,
               pd + nOffset + pParser->_nNaluUnitLength, nNaluLen);
        nOffset += (nNaluLen + pParser->_nNaluUnitLength);
        _nMediaLen += (4 + nNaluLen);
    }

    return 0;
}

int FlvParser::DumpH264(const std::string filepath) {
    std::fstream fout;
    fout.open(filepath.c_str(), std::ios_base::out | std::ios_base::binary);

    if (fout) {
        std::vector<Tag*>::iterator it_tag;
        it_tag = _vpTag.begin();

        for (it_tag; it_tag != _vpTag.end(); it_tag++) {
            if ((*it_tag)->_header.nTagType != 0x09) continue;

            fout.write((char*)(*it_tag)->_pMedia, (*it_tag)->_nMediaLen);
        }
    }

    fout.close();

    return 0;
}

int FlvParser::DumpAAC(const std::string filepath) {
    std::fstream fout;
    fout.open(filepath.c_str(), std::ios_base::out | std::ios_base::binary);

    if (fout) {
        std::vector<Tag*>::iterator it_tag;
        it_tag = _vpTag.begin();

        for (it_tag; it_tag != _vpTag.end(); it_tag++) {
            if ((*it_tag)->_header.nTagType != 0x08) continue;

            AudioTag* ad_tag = (AudioTag*)(*it_tag);
            if (ad_tag->_nSoundFormat != 10) continue;

            if (ad_tag->_nMediaLen == 0) continue;

            fout.write((char*)(*it_tag)->_pMedia, (*it_tag)->_nMediaLen);
        }
    }

    fout.close();

    return 0;
}

int FlvParser::DumpFlv(const std::string filepath) {}

void FlvParser::SetStat() {
    for (size_t i = 0; i < _vpTag.size(); i++) {
        switch (_vpTag[i]->_header.nTagType) {
            case 0x08:
                _sStat.nAudioNum++;
                break;
            case 0x09:
                SetVideoStat(_vpTag[i]);
                break;
            case 0x12:
                _sStat.nMetaNum++;
                break;
        }
    }
}

void FlvParser::SetVideoStat(Tag* pTag) {
    _sStat.nVideoNum++;
    _sStat.nMaxTimeStamp = pTag->_header.nTimeSamp;

    if (pTag->_pTagData[0] == 0x17 && pTag->_pTagData[1] == 0x00) {
        _sStat.nLengthSize = (pTag->_pTagData[9] & 0x03) + 1;
    }
}

int FlvParser::MetaTag::parseMeta(FlvParser* pParser) {
    uint8_t* pd = _pTagData;
    int dataSize = _header.nDataSize;

    uint32_t arrayLen = 0;
    uint32_t offset = 13;  // Type + Value_Size + Value占用13字节

    uint32_t nameLen = 0;
    double doubleValue = 0;
    std::string strValue = "";
    bool boolValue = false;
    uint32_t valueLen = 0;
    uint8_t u8Value = 0;
    if (pd[offset++] == 0x08)  // 0x8 onMetaData
    {
        arrayLen = SolveU32(pd + offset);
        offset += 4;  // 跳过 [ECMAArrayLength]占用的字节
        printf("ArrayLen = %d\n", arrayLen);
    } else {
        printf("metadata format error!!!");
        return -1;
    }

    for (uint32_t i = 0; i < arrayLen; i++) {
        doubleValue = 0;
        boolValue = false;
        strValue = "";
        // 读取字段长度
        nameLen = SolveU16(pd + offset);
        offset += 2;  // 跳过2字节字段长度

        char name[nameLen + 1];  // 获取字段名称

        memset(name, 0, sizeof(name));
        memcpy(name, &pd[offset], nameLen);
        name[nameLen + 1] = '\0';
        offset += nameLen;  // 跳过字段名占用的长度

        uint8_t amfType = pd[offset++];
        switch (amfType)  // 判别值的类型
        {
            case 0x0:  // Number type, 就是double类型, 占用8字节
                doubleValue = hexStr2Double(&pd[offset], 8);
                offset += 8;  // 跳过8字节
                break;

            case 0x1:  // Boolean type, 占用1字节
                u8Value = SolveU8(pd + offset);
                offset += 1;  // 跳过1字节
                if (u8Value != 0x00)
                    boolValue = true;
                else
                    boolValue = false;
                break;

            case 0x2:  // String type
                valueLen = SolveU16(pd + offset);
                offset += 2;  // 跳过2字节 length
                strValue.append(pd + offset, pd + offset + valueLen);
                strValue.append("");
                offset += valueLen;  // 跳过字段的值的长度
                break;

            default:
                printf("un handle amfType:%d\n", amfType);
                break;
        }

        if (strncmp(name, "duration", 8) == 0) {
            m_duration = doubleValue;
        } else if (strncmp(name, "width", 5) == 0) {
            m_width = doubleValue;
        } else if (strncmp(name, "height", 6) == 0) {
            m_height = doubleValue;
        } else if (strncmp(name, "videodatarate", 13) == 0) {
            m_videodatarate = doubleValue;
        } else if (strncmp(name, "framerate", 9) == 0) {
            m_framerate = doubleValue;
        } else if (strncmp(name, "videocodecid", 12) == 0) {
            m_videocodecid = doubleValue;
        } else if (strncmp(name, "audiodatarate", 13) == 0) {
            m_audiodatarate = doubleValue;
        } else if (strncmp(name, "audiosamplerate", 15) == 0) {
            m_audiosamplerate = doubleValue;
        } else if (strncmp(name, "audiosamplesize", 15) == 0) {
            m_audiosamplesize = doubleValue;
        } else if (strncmp(name, "stereo", 6) == 0) {
            m_stereo = boolValue;
        } else if (strncmp(name, "audiocodecid", 12) == 0) {
            m_audiocodecid = doubleValue;
        } else if (strncmp(name, "major_brand", 11) == 0) {
            m_major_brand = strValue;
        } else if (strncmp(name, "minor_version", 13) == 0) {
            m_minor_version = strValue;
        } else if (strncmp(name, "compatible_brands", 17) == 0) {
            m_compatible_brand = strValue;
        } else if (strncmp(name, "encoder", 7) == 0) {
            m_encoder = strValue;
        } else if (strncmp(name, "filesize", 8) == 0) {
            m_filesize = doubleValue;
        }
    }

    printMeta();
    return 1;
}

void FlvParser::MetaTag::printMeta() {
    printf("\nduration: %0.2lfs, filesize: %.0lfbytes\n", m_duration,
           m_filesize);

    printf("width: %0.0lf, height: %0.0lf\n", m_width, m_height);
    printf("videodatarate: %0.2lfkbps, framerate: %0.0lffps\n", m_videodatarate,
           m_framerate);
    printf("videocodecid: %0.0lf\n", m_videocodecid);

    printf("audiodatarate: %0.2lfkbps, audiosamplerate: %0.0lfKhz\n",
           m_audiodatarate, m_audiosamplerate);
    printf("audiosamplesize: %0.0lfbit, stereo: %d\n", m_audiosamplesize,
           m_stereo);
    printf("audiocodecid: %0.0lf\n", m_audiocodecid);

    printf("major_brand: %s, minor_version: %s\n", m_major_brand.c_str(),
           m_minor_version.c_str());
    printf("compatible_brands: %s, encoder: %s\n\n", m_compatible_brand.c_str(),
           m_encoder.c_str());
}

double FlvParser::MetaTag::hexStr2Double(const unsigned char* hex,
                                         const unsigned int length) {
    double ret = 0;
    char hexstr[length * 2];
    memset(hexstr, 0, length * 2);

    for (int i = 0; i < length; i++) {
        std::sprintf(hexstr + i * 2, "%02x", hex[i]);
    }

    sscanf(hexstr, "%llx", (unsigned long long*)&ret);

    return ret;
}

void FlvParser::PrintInfo() {
    SetStat();

    std::cout << "vnum:" << _sStat.nVideoNum << ",anum:" << _sStat.nAudioNum
              << std::endl;
    std::cout << "maxTimeStamp:" << _sStat.nMaxTimeStamp
              << ",nLengthSize:" << _sStat.nLengthSize << std::endl;
}