#include "VideoEncoder.h"
#include <stdexcept>

namespace PTAMM {

static const int FRAME_WIDTH = 640;
static const int FRAME_HEIGHT = 480;
static const char STREAM_FILENAME[] = "rtp://127.0.0.1:5004";

AVFrame *AllocPicture(PixelFormat pix_fmt, int width, int height)
{
  // Allocate a frame
  AVFrame* pic = avcodec_alloc_frame();
  if (!pic) {
    throw std::runtime_error("avcodec_alloc_frame failed");
  }

  if (av_image_alloc(pic->data, pic->linesize,
      width, height, pix_fmt, 1) < 0)
  {
    throw std::runtime_error("av_image_alloc failed");
  }

  pic->width = FRAME_WIDTH;
  pic->height = FRAME_HEIGHT;
  return pic;
}

static void print_sdp(AVFormatContext **avc, int n)
{
  char sdp[2048];
  avf_sdp_create(avc, n, sdp, sizeof(sdp));
  printf("SDP:\n%s\n", sdp);
  fflush(stdout);
}

VideoEncoder::VideoEncoder()
  : format_ctx_(0)
  , stream_(0)
  , pic_(0)
  , outbuf_(0)
  , outbuf_size_(0)
  , frame_count_(0)
{
  // Register all formats and codecs
  av_register_all();
  avcodec_register_all();

  // Init networking
  avformat_network_init();

  // Find format
  format_ = av_guess_format("mpegts", NULL, NULL);

  // allocate the AVFormatContext
  format_ctx_ = avformat_alloc_context();
  if (!format_ctx_) {
    throw std::runtime_error("avformat_alloc_context failed");
  }

  format_ctx_->oformat = format_;
  snprintf(format_ctx_->filename, sizeof(format_ctx_->filename), "%s", STREAM_FILENAME);

  stream_ = CreateVideoStream(format_ctx_);

  // Allocate encoder buffer
  outbuf_size_ = 4096*400; // Must be large enough to contain one encoded frame
  outbuf_ = new uint8_t[outbuf_size_];

  // Allocate AVIOContext
  int ret = avio_open(&io_ctx_, STREAM_FILENAME, AVIO_FLAG_WRITE);
  if (ret != 0) {
    throw std::runtime_error("avio_open failed");
  }

  format_ctx_->pb = io_ctx_;

  // Allocate a frame
  pic_ = AllocPicture(PIX_FMT_YUV420P, FRAME_WIDTH, FRAME_HEIGHT);

  // Print some debug info about the format
  av_dump_format(format_ctx_, 0, NULL, 1);

  // Begin the output by writing the container header
  avformat_write_header(format_ctx_, NULL);

  AVFormatContext* ac[] = { format_ctx_ };
  print_sdp(ac, 1);
}

VideoEncoder::~VideoEncoder()
{
  av_write_trailer(format_ctx_);
  avcodec_close(stream_->codec);

  for (size_t i = 0; i < format_ctx_->nb_streams; i++) {
    av_freep(&format_ctx_->streams[i]->codec);
    av_freep(&format_ctx_->streams[i]);
  }

  av_freep(&format_ctx_->pb->buffer);
  av_freep(&format_ctx_->pb);
  av_freep(&pic_->data[0]);
  av_freep(&pic_);

  delete[] outbuf_;
}

AVStream* VideoEncoder::CreateVideoStream(AVFormatContext *oc)
{
  AVStream* st = avformat_new_stream(oc, NULL);
  if (!st) {
    std::runtime_error("Could not alloc stream");
  }

  AVCodec* codec = avcodec_find_encoder(CODEC_ID_MPEG2VIDEO);
  if (!codec) {
    throw std::runtime_error("couldn't find mpeg2 encoder");
  }

  AVCodecContext* c = st->codec;

  avcodec_get_context_defaults3(c, codec);

  c->codec_id = CODEC_ID_MPEG2VIDEO;
  c->codec_type = AVMEDIA_TYPE_VIDEO;
//  c->bit_rate = 100000;
  c->width = FRAME_WIDTH;
  c->height = FRAME_HEIGHT;
  c->time_base.num = 1;
  c->time_base.den = 25;
  c->pix_fmt = PIX_FMT_YUV420P;

  AVDictionary* opts = NULL;
  if (avcodec_open2(c, codec, &opts) < 0) {
    throw std::runtime_error("avcodec_open failed");
  }

  return st;
}

// This function is not doing what it should do.
// The source image is in RGB format (24 bpp) and the dest frame is supposed
// to be YUV420P, so a conversion would be necessary
// I doubt this is the reason why the streaming is not working though.
void VideoEncoder::CopyFrameData(const uint8_t* src_frame, AVFrame* dest_frame)
{
  const uint8_t* src_pix = src_frame;
  uint8_t* dest_pix = dest_frame->data[0];

  for(int y = 0; y < dest_frame->height; ++y) {
    uint8_t* p = dest_pix;
    for(int x = 0; x < dest_frame->width; ++x) {
      *p++ = *src_pix++;
      *p++ = *src_pix++;
      *p++ = *src_pix++;
    }
    dest_pix += dest_frame->linesize[0];
  }
}

void VideoEncoder::EncodeFrame(const uint8_t* frame)
{
  AVCodecContext* c = stream_->codec;

  // Convert and load the frame data into the AVFrame struct
  CopyFrameData(frame, pic_);

  // Generate a PTS
  pic_->pts = (int64_t)(frame_count_ * av_q2d(c->time_base) * 1000);

  // Encode the image
  int out_size = avcodec_encode_video(c, outbuf_, outbuf_size_, pic_);
  // if zero size, it means the image was buffered
  if (out_size != 0) {
      AVPacket pkt;
      av_init_packet(&pkt);

      pkt.pts = pic_->pts;
      //pkt.stream_index = st->index;
      pkt.data = outbuf_;
      pkt.size = out_size;

      /* write the compressed frame in the media file */
      int ret = av_interleaved_write_frame(format_ctx_, &pkt);

      if (ret != 0) {
        throw std::runtime_error("Failed to write frame");
      }
  }

  ++frame_count_;
}


}