//          Copyright (c) Claus Hoefele.  All rights reserved.                //
//==============================================================================

#include "TheoraResourceWriter.h"
#include "../Image.h"

#include <system/assert.h>

#include <cstdio>
#include <stdint.h>

using namespace std;

//------------------------------------------------------------------------------
namespace
{
  const int VIDEO_FPS_NUMERATOR = 30;
  const int VIDEO_FPS_DENOMINATOR = 1;
  const int VIDEO_ASPECT_NUMERATOR = 0;
  const int VIDEO_ASPECT_DENOMINATOR = 0;
  const int VIDEO_RATE = 768000;
  const int VIDEO_QUALITY = 48;
}

//------------------------------------------------------------------------------
TheoraResourceWriter::TheoraResourceWriter(const string& fileName)
  : m_File(NULL)
  , m_TheoraState(NULL)
  , m_OggStreamState()
{
  const int width = 512;
  const int height = 512;

  // Theora has a divisible-by-sixteen restriction for the encoded 
  // frame size. This code scales the picture size up to the nearest 
  // multiple of 16 and calculates offsets to the real size.
  int frame_w = width + 15 & ~0xF;
  int frame_h = height + 15 & ~0xF;
  /*Force the offsets to be even so that chroma samples line up like we
     expect.*/
  int pic_x = (frame_w - width) >> 1 & ~1;
  int pic_y = (frame_h - height) >> 1 & ~1;

  th_info theoraInfo;
  th_info_init(&theoraInfo);
  theoraInfo.frame_width = frame_w;
  theoraInfo.frame_height = frame_h;
  theoraInfo.pic_width = width;
  theoraInfo.pic_height = height;
  theoraInfo.pic_x = pic_x;
  theoraInfo.pic_y = pic_y;
  theoraInfo.fps_numerator = VIDEO_FPS_NUMERATOR;
  theoraInfo.fps_denominator = VIDEO_FPS_DENOMINATOR;
  theoraInfo.aspect_numerator = VIDEO_ASPECT_NUMERATOR;
  theoraInfo.aspect_denominator = VIDEO_ASPECT_DENOMINATOR;
  theoraInfo.colorspace = TH_CS_UNSPECIFIED;
  theoraInfo.target_bitrate = VIDEO_RATE;
  theoraInfo.quality = VIDEO_QUALITY;
  theoraInfo.keyframe_granule_shift = 6;
  theoraInfo.pixel_fmt = TH_PF_420;

  m_TheoraState = th_encode_alloc(&theoraInfo);
  th_info_clear(&theoraInfo);

  // Init codec and open output file. 
  int result = open(fileName);
  ASSERT_M(result == 0, "Can't open libtheora.");
}

//------------------------------------------------------------------------------
TheoraResourceWriter::~TheoraResourceWriter()
{
  close();
}

//------------------------------------------------------------------------------
int TheoraResourceWriter::open(const string& fileName)
{
  m_File = fopen(fileName.c_str(), "wb");
  ASSERT(m_File);

  /* write the bitstream header packets with proper page interleave */
  th_comment theoraComment;
  th_comment_init(&theoraComment);

  ogg_packet oggPacket;
  if(th_encode_flushheader(m_TheoraState, &theoraComment, &oggPacket) <= 0)
  {
    fprintf(stderr, "Internal Theora library error.\n");
    return 1;
  }

  int errorCode = ogg_stream_init(&m_OggStreamState, rand());
  ASSERT(errorCode == 0);
  errorCode = ogg_stream_packetin(&m_OggStreamState, &oggPacket);
  ASSERT(errorCode == 0);
  
  ogg_page oggPage;
  if(ogg_stream_pageout(&m_OggStreamState, &oggPage) != 1)
  {
    fprintf(stderr, "Internal Ogg library error.\n");
    exit(1);
  }
  fwrite(oggPage.header, 1, oggPage.header_len, m_File);
  fwrite(oggPage.body, 1, oggPage.body_len, m_File);

  /* create the remaining theora headers */
  for(;;)
  {
    int errorCode = th_encode_flushheader(m_TheoraState, &theoraComment, &oggPacket);
    ASSERT_M(errorCode >= 0, "Encoder error.");
    if (errorCode <= 0)
    {
      break;
    }
    errorCode = ogg_stream_packetin(&m_OggStreamState, &oggPacket);
    ASSERT(errorCode == 0);
  }
  th_comment_clear(&theoraComment);

  /* Flush the rest of our headers. This ensures
     the actual data in each stream will start
     on a new page, as per spec. */
  for(;;)
  {
    int result = ogg_stream_flush(&m_OggStreamState, &oggPage);
    if(result < 0)
    {
      /* can't get here */
      fprintf(stderr, "Internal Ogg library error.\n");
      return 1;
    }

    if(result==0)
    {
      break;
    }

    fwrite(oggPage.header, 1, oggPage.header_len, m_File);
    fwrite(oggPage.body, 1, oggPage.body_len, m_File);
  }

  return 0;
}

//------------------------------------------------------------------------------
void TheoraResourceWriter::close()
{
  if ((m_File != NULL) && (m_TheoraState != NULL))
  {
    // Flush pending data.
    ogg_packet oggPacket;
    ogg_page oggPage;

    int errorCode = th_encode_packetout(m_TheoraState, 1, &oggPacket);  // flag for last packet set
    ASSERT(errorCode >= 0);
    errorCode = ogg_stream_pageout(&m_OggStreamState, &oggPage);
    ASSERT(errorCode >= 0);
    while (errorCode) 
    {
      fwrite(oggPage.header, oggPage.header_len, 1, m_File);
      fwrite(oggPage.body, oggPage.body_len, 1, m_File);
      errorCode = ogg_stream_pageout(&m_OggStreamState, &oggPage);
      ASSERT(errorCode >= 0);
    }
    
    errorCode = ogg_stream_flush(&m_OggStreamState, &oggPage);
    ASSERT(errorCode >= 0);
    if(errorCode)
    {
      fwrite(oggPage.header, oggPage.header_len, 1, m_File);
      fwrite(oggPage.body, oggPage.body_len, 1, m_File);
      errorCode = ogg_stream_flush(&m_OggStreamState, &oggPage);
      ASSERT(errorCode >= 0);
    }
  }

  if (m_File != NULL) 
  {
    fclose(m_File);
  }
  
  int errorCode = ogg_stream_clear(&m_OggStreamState);
  ASSERT(errorCode == 0);
}

//------------------------------------------------------------------------------
bool TheoraResourceWriter::write(/*const Image& image,*/ bool last)
{
  //const unsigned long width = static_cast<unsigned long>(image.getWidth());
  //const unsigned long height = static_cast<unsigned long>(image.getHeight());
  const int width = 512;
  const int height = 512;
  //const uint8_t* data = image.getData();

  int dst_c_dec_h = 2;
  int dst_c_dec_v = 2;
  int frame_w = width + 15 & ~0xF;
  int frame_h = height + 15 & ~0xF;
  int pic_x = (frame_w - width) >> 1 & ~1;
  int pic_y = (frame_h - height) >> 1 & ~1;

  int pic_sz = width * height;
  int frame_c_w = frame_w / dst_c_dec_h;
  int frame_c_h = frame_h / dst_c_dec_v;
  int c_w = (width + dst_c_dec_h - 1) / dst_c_dec_h;
  int c_h = (height + dst_c_dec_v - 1) / dst_c_dec_v;
  int c_sz = c_w * c_h;

  size_t y4m_dst_buf_sz = width * height + 
    2 * ((width + dst_c_dec_h - 1) / dst_c_dec_h) *
    ((height + dst_c_dec_v - 1) / dst_c_dec_v);

  unsigned char* yuvframe = (unsigned char *) malloc(y4m_dst_buf_sz);
  memset(yuvframe, 200, y4m_dst_buf_sz);

  th_ycbcr_buffer ycbcr;
  ycbcr[0].width = frame_w;
  ycbcr[0].height = frame_h;
  ycbcr[0].stride = width;
  ycbcr[0].data = yuvframe - pic_x - pic_y * width;
  ycbcr[1].width = frame_c_w;
  ycbcr[1].height = frame_c_h;
  ycbcr[1].stride = c_w;
  ycbcr[1].data = ycbcr[0].data + pic_sz - (pic_x / dst_c_dec_h) - 
    (pic_y / dst_c_dec_v) * c_w;
  ycbcr[2].width = frame_c_w;
  ycbcr[2].height = frame_c_h;
  ycbcr[2].stride = c_w;
  ycbcr[2].data = ycbcr[1].data+c_sz;
  int errorCode = th_encode_ycbcr_in(m_TheoraState, ycbcr);
  ASSERT_M(errorCode == 0, "Can't encode frame.");

  ogg_packet oggPacket;
  errorCode = th_encode_packetout(m_TheoraState, last ? 1 : 0, &oggPacket);
  ASSERT(errorCode >= 0);
  while(errorCode)
  {
    errorCode = ogg_stream_packetin(&m_OggStreamState, &oggPacket);
    ASSERT(errorCode == 0);
    errorCode = th_encode_packetout(m_TheoraState, last ? 1 : 0, &oggPacket);
    ASSERT(errorCode >= 0);
  }

  ogg_page oggPage;
  //int pageReady = ogg_stream_flush(&m_OggStreamState, &oggPage);
  int pageReady = ogg_stream_pageout(&m_OggStreamState, &oggPage);
  ASSERT(pageReady >= 0);
  
  if (pageReady)
  {
    fwrite(oggPage.header, 1, oggPage.header_len, m_File);
    fwrite(oggPage.body, 1, oggPage.body_len, m_File);
  }

  free(yuvframe);

  return true;
}