/*

 * H.265 video codec.

 * Copyright (c) 2013-2014 struktur AG, Dirk Farin <farin@struktur.de>

 *

 * This file is part of libde265.

 *

 * libde265 is free software: you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as

 * published by the Free Software Foundation, either version 3 of

 * the License, or (at your option) any later version.

 *

 * libde265 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 Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with libde265.  If not, see <http://www.gnu.org/licenses/>.

 */

#include "nal-parser.h"

#include <string.h>

#include <assert.h>

#include <stdlib.h>

#include <stdio.h>

#ifdef HAVE_CONFIG_H

#include "config.h"

#endif

NAL_unit::NAL_unit()

  : skipped_bytes(DE265_SKIPPED_BYTES_INITIAL_SIZE)

{

  pts=0;

  user_data = NULL;

  nal_data = NULL;

  data_size = 0;

  capacity = 0;

}

NAL_unit::~NAL_unit()

{

  free(nal_data);

}

void NAL_unit::clear()

{

  header = nal_header();

  pts = 0;

  user_data = NULL;

  // set size to zero but keep memory

  data_size = 0;

  skipped_bytes.clear();

}

LIBDE265_CHECK_RESULT bool NAL_unit::resize(int new_size)

{

  if (capacity < new_size) {

    unsigned char* newbuffer = (unsigned char*)malloc(new_size);

    if (newbuffer == NULL) {

      return false;

    }

    if (nal_data != NULL) {

      memcpy(newbuffer, nal_data, data_size);

      free(nal_data);

    }

    nal_data = newbuffer;

    capacity = new_size;

  }

  return true;

}

LIBDE265_CHECK_RESULT bool NAL_unit::append(const unsigned char* in_data, int n)

{

  if (!resize(data_size + n)) {

    return false;

  }

  memcpy(nal_data + data_size, in_data, n);

  data_size += n;

  return true;

}

bool LIBDE265_CHECK_RESULT NAL_unit::set_data(const unsigned char* in_data, int n)

{

  if (!resize(n)) {

    return false;

  }

  memcpy(nal_data, in_data, n);

  data_size = n;

  return true;

}

void NAL_unit::insert_skipped_byte(int pos)

{

  skipped_bytes.push_back(pos);

}

int NAL_unit::num_skipped_bytes_before(int byte_position, int headerLength) const

{

  for (int k=skipped_bytes.size()-1;k>=0;k--)

    if (skipped_bytes[k]-headerLength <= byte_position) {

      return k+1;

    }

  return 0;

}

void NAL_unit::remove_stuffing_bytes()

{

  uint8_t* p = data();

  for (int i=0;i<size()-2;i++)

    {

#if 0

        for (int k=i;k<i+64;k++) 

          if (i*0+k<size()) {

            printf("%c%02x", (k==i) ? '[':' ', data()[k]);

          }

        printf("\n");

#endif

      if (p[2]!=3 && p[2]!=0) {

        // fast forward 3 bytes (2+1)

        p+=2;

        i+=2;

      }

      else {

        if (p[0]==0 && p[1]==0 && p[2]==3) {

          //printf("SKIP NAL @ %d\n",i+2+num_skipped_bytes);

          insert_skipped_byte(i+2 + num_skipped_bytes());

          memmove(p+2, p+3, size()-i-3);

          set_size(size()-1);

          p++;

          i++;

        }

      }

      p++;

    }

}

NAL_Parser::NAL_Parser()

{

  end_of_stream = false;

  end_of_frame = false;

  input_push_state = 0;

  pending_input_NAL = NULL;

  nBytes_in_NAL_queue = 0;

}

NAL_Parser::~NAL_Parser()

{

  // --- free NAL queues ---

  // empty NAL queue

  NAL_unit* nal;

  while ( (nal = pop_from_NAL_queue()) ) {

    free_NAL_unit(nal);

  }

  // free the pending input NAL

  if (pending_input_NAL != NULL) {

    free_NAL_unit(pending_input_NAL);

  }

  // free all NALs in free-list

  for (int i=0;i<NAL_free_list.size();i++) {

    delete NAL_free_list[i];

  }

}

LIBDE265_CHECK_RESULT NAL_unit* NAL_Parser::alloc_NAL_unit(int size)

{

  NAL_unit* nal;

  // --- get NAL-unit object ---

  if (NAL_free_list.size() > 0) {

    nal = NAL_free_list.back();

    NAL_free_list.pop_back();

  }

  else {

    nal = new NAL_unit;

  }

  nal->clear();

  if (!nal->resize(size)) {

    free_NAL_unit(nal);

    return NULL;

  }

  return nal;

}

void NAL_Parser::free_NAL_unit(NAL_unit* nal)

{

  if (nal == NULL) {

    // Allow calling with NULL just like regular "free()"

    return;

  }

  if (NAL_free_list.size() < DE265_NAL_FREE_LIST_SIZE) {

    NAL_free_list.push_back(nal);

  }

  else {

    delete nal;

  }

}

NAL_unit* NAL_Parser::pop_from_NAL_queue()

{

  if (NAL_queue.empty()) {

    return NULL;

  }

  else {

    NAL_unit* nal = NAL_queue.front();

    NAL_queue.pop();

    nBytes_in_NAL_queue -= nal->size();

    return nal;

  }

}

void NAL_Parser::push_to_NAL_queue(NAL_unit* nal)

{

  NAL_queue.push(nal);

  nBytes_in_NAL_queue += nal->size();

}

de265_error NAL_Parser::push_data(const unsigned char* data, int len,

                                  de265_PTS pts, void* user_data)

{

  end_of_frame = false;

  if (pending_input_NAL == NULL) {

    pending_input_NAL = alloc_NAL_unit(len+3);

    if (pending_input_NAL == NULL) {

      return DE265_ERROR_OUT_OF_MEMORY;

    }

    pending_input_NAL->pts = pts;

    pending_input_NAL->user_data = user_data;

  }

  NAL_unit* nal = pending_input_NAL; // shortcut

  // Resize output buffer so that complete input would fit.

  // We add 3, because in the worst case 3 extra bytes are created for an input byte.

  if (!nal->resize(nal->size() + len + 3)) {

    return DE265_ERROR_OUT_OF_MEMORY;

  }

  unsigned char* out = nal->data() + nal->size();

  for (int i=0;i<len;i++) {

    /*

    printf("state=%d input=%02x (%p) (output size: %d)\n",ctx->input_push_state, *data, data,

           out - ctx->nal_data.data);

    */

    switch (input_push_state) {

    case 0:

    case 1:

      if (*data == 0) { input_push_state++; }

      else { input_push_state=0; }

      break;

    case 2:

      if      (*data == 1) { input_push_state=3; } // nal->clear_skipped_bytes(); }

      else if (*data == 0) { } // *out++ = 0; }

      else { input_push_state=0; }

      break;

    case 3:

      *out++ = *data;

      input_push_state = 4;

      break;

    case 4:

      *out++ = *data;

      input_push_state = 5;

      break;

    case 5:

      if (*data==0) { input_push_state=6; }

      else { *out++ = *data; }

      break;

    case 6:

      if (*data==0) { input_push_state=7; }

      else {

        *out++ = 0;

        *out++ = *data;

        input_push_state=5;

      }

      break;

    case 7:

      if      (*data==0) { *out++ = 0; }

      else if (*data==3) {

        *out++ = 0; *out++ = 0; input_push_state=5;

        // remember which byte we removed

        nal->insert_skipped_byte((out - nal->data()) + nal->num_skipped_bytes());

      }

      else if (*data==1) {

#if DEBUG_INSERT_STREAM_ERRORS

        if ((rand()%100)<90 && nal_data.size>0) {

          int pos = rand()%nal_data.size;

          int bit = rand()%8;

          nal->nal_data.data[pos] ^= 1<<bit;

          //printf("inserted error...\n");

        }

#endif

        nal->set_size(out - nal->data());;

        // push this NAL decoder queue

        push_to_NAL_queue(nal);

        // initialize new, empty NAL unit

        pending_input_NAL = alloc_NAL_unit(len+3);

        if (pending_input_NAL == NULL) {

          return DE265_ERROR_OUT_OF_MEMORY;

        }

        pending_input_NAL->pts = pts;

        pending_input_NAL->user_data = user_data;

        nal = pending_input_NAL;

        out = nal->data();

        input_push_state=3;

        //nal->clear_skipped_bytes();

      }

      else {

        *out++ = 0;

        *out++ = 0;

        *out++ = *data;

        input_push_state=5;

      }

      break;

    }

    data++;

  }

  nal->set_size(out - nal->data());

  return DE265_OK;

}

de265_error NAL_Parser::push_NAL(const unsigned char* data, int len,

                                 de265_PTS pts, void* user_data)

{

  // Cannot use byte-stream input and NAL input at the same time.

  assert(pending_input_NAL == NULL);

  end_of_frame = false;

  NAL_unit* nal = alloc_NAL_unit(len);

  if (nal == NULL || !nal->set_data(data, len)) {

    free_NAL_unit(nal);

    return DE265_ERROR_OUT_OF_MEMORY;

  }

  nal->pts = pts;

  nal->user_data = user_data;

  nal->remove_stuffing_bytes();

  push_to_NAL_queue(nal);

  return DE265_OK;

}

de265_error NAL_Parser::flush_data()

{

  if (pending_input_NAL) {

    NAL_unit* nal = pending_input_NAL;

    uint8_t null[2] = { 0,0 };

    // append bytes that are implied by the push state

    if (input_push_state==6) {

      if (!nal->append(null,1)) {

        return DE265_ERROR_OUT_OF_MEMORY;

      }

    }

    if (input_push_state==7) {

      if (!nal->append(null,2)) {

        return DE265_ERROR_OUT_OF_MEMORY;

      }

    }

    // only push the NAL if it contains at least the NAL header

    if (input_push_state>=5) {

      push_to_NAL_queue(nal);

      pending_input_NAL = NULL;

    }

    input_push_state = 0;

  }

  return DE265_OK;

}

void NAL_Parser::remove_pending_input_data()

{

  // --- remove pending input data ---

  if (pending_input_NAL) {

    free_NAL_unit(pending_input_NAL);

    pending_input_NAL = NULL;

  }

  for (;;) {

    NAL_unit* nal = pop_from_NAL_queue();

    if (nal) { free_NAL_unit(nal); }

    else break;

  }

  input_push_state = 0;

  nBytes_in_NAL_queue = 0;

}