/* 
 * Copyright (C) 2020 Teluu Inc. (http://www.teluu.com)
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA 
 */
#include <pjmedia-codec/vpx_packetizer.h>
#include <pjmedia/errno.h>
#include <pjmedia/types.h>
#include <pjmedia/vid_codec_util.h>
#include <pj/assert.h>
#include <pj/errno.h>
#include <pj/log.h>
#include <pj/pool.h>
#include <pj/string.h>


#if defined(PJMEDIA_HAS_VIDEO) && (PJMEDIA_HAS_VIDEO != 0)

#define THIS_FILE		"vpx_packetizer.c"

/* VPX packetizer definition */
struct pjmedia_vpx_packetizer
{
    /* Current settings */
    pjmedia_vpx_packetizer_cfg cfg;
};

/*
 * Initialize VPX packetizer.
 */
PJ_DEF(void) pjmedia_vpx_packetizer_cfg_default(pjmedia_vpx_packetizer_cfg *cfg)
{
    pj_bzero(cfg, sizeof(*cfg));

    cfg->fmt_id = PJMEDIA_FORMAT_VP8;
    cfg->mtu =PJMEDIA_MAX_VID_PAYLOAD_SIZE;
}

/*
 * Create vpx packetizer.
 */
PJ_DEF(pj_status_t) pjmedia_vpx_packetizer_create(
				pj_pool_t *pool,
				const pjmedia_vpx_packetizer_cfg *cfg,
				pjmedia_vpx_packetizer **p)
{
    pjmedia_vpx_packetizer *p_;

    PJ_ASSERT_RETURN(pool && p, PJ_EINVAL);

    if (cfg && cfg->fmt_id != PJMEDIA_FORMAT_VP8 &&
	cfg->fmt_id != PJMEDIA_FORMAT_VP9)
    {
	return PJ_ENOTSUP;
    }

    p_ = PJ_POOL_ZALLOC_T(pool, pjmedia_vpx_packetizer);
    if (cfg) {
	pj_memcpy(&p_->cfg, cfg, sizeof(*cfg));
    } else {
	pjmedia_vpx_packetizer_cfg_default(&p_->cfg);
    }
    *p = p_;

    return PJ_SUCCESS;
}

/*
 * Generate an RTP payload from H.264 frame bitstream, in-place processing.
 */
PJ_DEF(pj_status_t) pjmedia_vpx_packetize(const pjmedia_vpx_packetizer *pktz,
					  pj_size_t bits_len,
                                          unsigned *bits_pos,
                                          pj_bool_t is_keyframe,
                                          pj_uint8_t **payload,
                                          pj_size_t *payload_len)
{
    unsigned payload_desc_size = 1;
    unsigned max_size = pktz->cfg.mtu - payload_desc_size;
    unsigned remaining_size = bits_len - *bits_pos;
    unsigned out_size = *payload_len;
    pj_uint8_t *bits = *payload;

    *payload_len = PJ_MIN(remaining_size, max_size);
    if (*payload_len + payload_desc_size > out_size)
	return PJMEDIA_CODEC_EFRMTOOSHORT;

    /* Set payload header */
    bits[0] = 0;
    if (pktz->cfg.fmt_id == PJMEDIA_FORMAT_VP8) {
	/* Set N: Non-reference frame */
        if (!is_keyframe) bits[0] |= 0x20;
        /* Set S: Start of VP8 partition. */
        if (*bits_pos == 0) bits[0] |= 0x10;
    } else if (pktz->cfg.fmt_id == PJMEDIA_FORMAT_VP9) {
	/* Set P: Inter-picture predicted frame */
        if (!is_keyframe) bits[0] |= 0x40;
        /* Set B: Start of a frame */
        if (*bits_pos == 0) bits[0] |= 0x8;
        /* Set E: End of a frame */
        if (*bits_pos + *payload_len == bits_len) {
            bits[0] |= 0x4;
	}
    }
    return PJ_SUCCESS;
}


/*
 * Append RTP payload to a VPX picture bitstream
 */
PJ_DEF(pj_status_t) pjmedia_vpx_unpacketize(pjmedia_vpx_packetizer *pktz,
					    const pj_uint8_t *payload,
                                            pj_size_t payload_len,
					    unsigned  *payload_desc_len)
{
    unsigned desc_len = 1;
    pj_uint8_t *p = (pj_uint8_t *)payload;

#define INC_DESC_LEN() {if (++desc_len >= payload_len) return PJ_ETOOSMALL;}

    if (payload_len <= desc_len) return PJ_ETOOSMALL;

    if (pktz->cfg.fmt_id == PJMEDIA_FORMAT_VP8) {
        /*  0 1 2 3 4 5 6 7
         * +-+-+-+-+-+-+-+-+
         * |X|R|N|S|R| PID | (REQUIRED)
         */
	/* X: Extended control bits present. */
	if (p[0] & 0x80) {
	    INC_DESC_LEN();
	    /* |I|L|T|K| RSV   | */
	    /* I: PictureID present. */
	    if (p[1] & 0x80) {
	    	INC_DESC_LEN();
	    	/* If M bit is set, the PID field MUST contain 15 bits. */
	    	if (p[2] & 0x80) INC_DESC_LEN();
	    }
	    /* L: TL0PICIDX present. */
	    if (p[1] & 0x40) INC_DESC_LEN();
	    /* T: TID present or K: KEYIDX present. */
	    if ((p[1] & 0x20) || (p[1] & 0x10)) INC_DESC_LEN();
	}

    } else if (pktz->cfg.fmt_id == PJMEDIA_FORMAT_VP9) {
        /*  0 1 2 3 4 5 6 7
         * +-+-+-+-+-+-+-+-+
         * |I|P|L|F|B|E|V|-| (REQUIRED)
         */
        /* I: Picture ID (PID) present. */
	if (p[0] & 0x80) {
	    INC_DESC_LEN();
	    /* If M bit is set, the PID field MUST contain 15 bits. */
	    if (p[1] & 0x80) INC_DESC_LEN();
	}
	/* L: Layer indices present. */
	if (p[0] & 0x20) {
	    INC_DESC_LEN();
	    if (!(p[0] & 0x10)) INC_DESC_LEN();
	}
	/* F: Flexible mode.
	 * I must also be set to 1, and if P is set, there's up to 3
	 * reference index.
	 */
	if ((p[0] & 0x10) && (p[0] & 0x80) && (p[0] & 0x40)) {
	    unsigned char *q = p + desc_len;

	    INC_DESC_LEN();
	    if (*q & 0x1) {
	    	q++;
	    	INC_DESC_LEN();
	    	if (*q & 0x1) {
	    	    q++;
	    	    INC_DESC_LEN();
	    	}
	    }
	}
	/* V: Scalability structure (SS) data present. */
	if (p[0] & 0x2) {
	    unsigned char *q = p + desc_len;
	    unsigned N_S = (*q >> 5) + 1;

	    INC_DESC_LEN();
	    /* Y: Each spatial layer's frame resolution present. */
	    if (*q & 0x10) desc_len += N_S * 4;

	    /* G: PG description present flag. */
	    if (*q & 0x8) {
	    	unsigned j;
	    	unsigned N_G = *(p + desc_len);

	    	INC_DESC_LEN();
	    	for (j = 0; j< N_G; j++) {
	    	    unsigned R;

	    	    q = p + desc_len;
	    	    INC_DESC_LEN();
	    	    R = (*q & 0x0F) >> 2;
	    	    desc_len += R;
	    	    if (desc_len >= payload_len)
	    	    	return PJ_ETOOSMALL;
	    	}
	    }
	}
    }
#undef INC_DESC_LEN

    *payload_desc_len = desc_len;
    return PJ_SUCCESS;
}


#endif /* PJMEDIA_HAS_VIDEO */
