/*
 * Copyright (C) 2020 Rockchip Electronics Co. LTD
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#undef  ROCKCHIP_LOG_TAG
#define ROCKCHIP_LOG_TAG    "C2RKNaluParser"

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

#include "rk_mpi.h"
#include "C2RKNaluParser.h"
#include "C2RKLog.h"


#define H264_NALU_TYPE_SPS              7
#define H264_PROFILE_IDC_HIGH10       110
#define H265_NALU_TYPE_SPS             33
#define H265_MAX_VPS_COUNT             16
#define H265_MAX_SUB_LAYERS             7
#define H265_PROFILE_IDC_MAIN_10        2


int32_t C2RKNaluParser::detechAVCBitDepth(uint8_t *buf, int32_t size) {
    BitReadContext  gbCtx;
    BitReadContext *gb = &gbCtx;
    int32_t startCodeLen = 0;
    int32_t value = 0;
    int32_t bitDepth = 8;

    c2_set_bitread_ctx(gb, buf, size);
    c2_set_pre_detection(gb);
    if (!c2_update_curbyte(gb)) {
        c2_err("failed to update curbyte, skipping.");
        goto error;
    }

    /*
     * ExtraData carry h264 sps_info in two ways.
     * 1. start with 0x000001 or 0x00000001
     * 2. AVC extraData configuration
     */

    if (buf[0] == 0x00 && buf[1] == 0x00 && buf[2] == 0x01) {
        startCodeLen = 3;
    } else if (buf[0] == 0x00 && buf[1] == 0x00 && buf[2] == 0x00 && buf[3] == 0x01) {
        startCodeLen = 4;
    } else {
        startCodeLen = 0;
    }

    if (startCodeLen > 0) {
        SKIP_BITS(gb, startCodeLen * 8);
    } else {
        // AVC extraData configuration
        SKIP_BITS(gb, 32);
        SKIP_BITS(gb, 16);

        SKIP_BITS(gb, 16);  // sequenceParameterSetLength
    }

    /* parse h264 sps info */
    READ_ONEBIT(gb, &value);  // forbidden_bit

    SKIP_BITS(gb, 2);  // nal_ref_idc
    READ_BITS(gb, 5, &value);  // nalu_type
    // stop traversal if not SPS nalu type
    if (value != H264_NALU_TYPE_SPS) {
        goto error;
    }

    READ_BITS(gb, 8, &value);  // profile_idc
    if (value == H264_PROFILE_IDC_HIGH10) {
        bitDepth = 10;
    } else {
        bitDepth = 8;
    }

    return bitDepth;

__BR_ERR:
error:
    bitDepth = 8;
    return bitDepth;
}

bool C2RKNaluParser::detachHEVCNalSPS(BitReadContext *gb, int32_t *bitDepth) {
    int32_t value = 0;

    READ_BITS(gb, 4, &value); // vps-id
    if (value > H265_MAX_VPS_COUNT) {
        c2_err("VPS id out of range: %d", value);
        goto error;
    }

    READ_BITS(gb, 3, &value);
    value += 1;

    if (value > H265_MAX_SUB_LAYERS) {
        c2_err("sps_max_sub_layers out of range: %d", value);
        goto error;
    }

    SKIP_BITS(gb, 1); // temporal_id_nesting_flag

    SKIP_BITS(gb, 3); // profile_space & tier_flag
    READ_BITS(gb, 5, &value); // profile_idc

    if (value == H265_PROFILE_IDC_MAIN_10) {
        *bitDepth = 10;
    } else {
        *bitDepth = 8;
    }

    return true;

__BR_ERR:
error:
    return false;
}

bool C2RKNaluParser::detachHEVCNalUnit(uint8_t *buf, int32_t size, int32_t *bitDepth) {
    BitReadContext gb_ctx;
    BitReadContext *gb = &gb_ctx;
    int32_t nalUnitType = 0;
    int32_t nuhLayerId = 0;
    int32_t temporalId = 0;

    c2_set_bitread_ctx(gb, buf, size);
    c2_set_pre_detection(gb);
    if (!c2_update_curbyte(gb)) {
        c2_err("failed to update curbyte, skipping.");
        return false;
    }

    SKIP_BITS(gb, 1); /* this bit should be zero */
    READ_BITS(gb, 6, &nalUnitType);
    READ_BITS(gb, 6, &nuhLayerId);
    READ_BITS(gb, 3, &temporalId);

    temporalId = temporalId -1;

    c2_trace("nal_unit_type: %d, nuh_layer_id: %d temporal_id: %d",
             nalUnitType, nuhLayerId, temporalId);

    if (temporalId < 0) {
        c2_err("Invalid NAL unit %d, skipping.", nalUnitType);
        goto error;
    }

    if (nalUnitType == H265_NALU_TYPE_SPS) {
        if (detachHEVCNalSPS(gb, bitDepth)) {
            return true;
        }
    }

__BR_ERR:
error:
    return false;
}

int32_t C2RKNaluParser::detechHEVCBitDepth(uint8_t *buf, int32_t size) {
    int32_t bitDepth = 8;

    if (buf[0] || buf[1] || buf[2] > 1) {
        int32_t i = 0, j = 0;
        int32_t nalLenSize = 0;
        uint32_t numOfArrays = 0, numOfNals = 0;

        /* It seems the extradata is encoded as hvcC format.
         * Temporarily, we support configurationVersion==0 until 14496-15 3rd
         * is finalized. When finalized, configurationVersion will be 1 and we
         * can recognize hvcC by checking if h265dctx->extradata[0]==1 or not. */
        if (size < 7) {
            goto error;
        }

        c2_info("extradata is encoded as hvcC format");

        nalLenSize = 1 + (buf[14 + 7] & 3);
        buf += 22;
        size -= 22;
        numOfArrays = static_cast<char>(buf[0]);
        buf += 1;
        size -= 1;
        for (i = 0; i < numOfArrays; i++) {
            buf += 1;
            size -= 1;
            // Num of nals
            numOfNals = buf[0] << 8 | buf[1];
            buf += 2;
            size -= 2;

            for (j = 0; j < numOfNals; j++) {
                uint32_t length = 0;
                if (size < 2) {
                    goto error;
                }

                length = buf[0] << 8 | buf[1];

                buf += 2;
                size -= 2;
                if (size < length) {
                    goto error;
                }
                if (detachHEVCNalUnit(buf, length, &bitDepth)) {
                    return bitDepth;
                }
                buf += length;
                size -= length;
            }
        }
    } else {
        int32_t i;
        for (i = 0; i < size - 4; i++) {
            // find sps start code
            if (buf[i] == 0x00 && buf[i + 1] == 0x00 &&
                buf[i + 2] == 0x01 && ((buf[i + 3] & 0x7f) >> 1) == H265_NALU_TYPE_SPS) {
                c2_info("find h265 sps start code.");
                i += 3;
                if (detachHEVCNalUnit(buf + i, size - i, &bitDepth)) {
                    return bitDepth;
                }
            }
        }
    }

error:
    bitDepth = 8;
    return bitDepth;
}

int32_t C2RKNaluParser::detectBitDepth(uint8_t *buf, int32_t size, int32_t coding) {
    int32_t bitDepth = 0;

    switch (coding) {
        case MPP_VIDEO_CodingAVC: {
            bitDepth = detechAVCBitDepth(buf, size);
            break;
        }
        case MPP_VIDEO_CodingHEVC: {
            bitDepth = detechHEVCBitDepth(buf, size);
            break;
        }
        default: {
            bitDepth = 8;
            c2_trace("not support coding %d, set default 8bit", coding);
            break;
        }
    }
    return bitDepth;
}