#ifndef FFMPEGDECODER_H
#define FFMPEGDECODER_H

#include <iostream>
#include <memory>
#include <QImage>
#include <QMutexLocker>

extern "C"{
#include <libavcodec/avcodec.h>
#include <libavutil/channel_layout.h>
#include <libavutil/common.h>
#include <libavutil/frame.h>
#include <libavutil/samplefmt.h>
#include <libavutil/opt.h>
#include <libavutil/imgutils.h>
#include <libavutil/parseutils.h>
#include <libavutil/mem.h>
#include <libswscale/swscale.h>
#include <libavformat/avformat.h>
}

class FFmpegDecoder
{
public:
  FFmpegDecoder(AVCodecID codecID=AV_CODEC_ID_H264);
  ~FFmpegDecoder();

public:
  bool init(AVCodecID codecID);
  void parsePkt(uint8_t* data, int len);
  void parsePkt(AVPacket &packet);
  QImage getImage() const;
  void clear();

private:
  void decode();
  void setImage(const QImage &newImage);

private:
  const AVCodec                 *videoCodec   = nullptr;
  AVCodecContext                *videoCodecCtx= nullptr;
  AVCodecParserContext          *parser       = nullptr;
  AVPacket                      *pkt          = nullptr;
  AVFrame                       *yuvFrame     = nullptr;
  AVFrame                       *rgbFrame     = nullptr;
  SwsContext                    *img_ctx      = nullptr;
  unsigned char                 *out_buffer   = nullptr;
  AVCodecID                     id            = AV_CODEC_ID_H264;
  QMutex                        mutex;
  QImage                        image;
  bool                          openCodec     = false;
};

using FFmpegDecoderPtr = std::shared_ptr<FFmpegDecoder>;

#endif // FFMPEGDECODER_H

#include "FFmpegDecoder.h"
#include <QDebug>

FFmpegDecoder::FFmpegDecoder(AVCodecID codecID)
{
  id = codecID;
  pkt = av_packet_alloc();
  yuvFrame = av_frame_alloc();
  rgbFrame = av_frame_alloc();
}

FFmpegDecoder::~FFmpegDecoder()
{
  if(!parser) av_parser_close(parser);
  if(!pkt) av_packet_free(&pkt);
  if(!yuvFrame) av_frame_free(&yuvFrame);
  if(!rgbFrame) av_frame_free(&rgbFrame);
  if(!videoCodecCtx) avcodec_free_context(&videoCodecCtx);
  if(!videoCodecCtx) avcodec_close(videoCodecCtx);
}

bool FFmpegDecoder::init(AVCodecID codecID)
{
  id = codecID;
  if(!videoCodec){
    if(!(videoCodec = avcodec_find_decoder(id))){
      printf("Cannot find valid decode codec.\n");
      return false;
    }
  }
  if(!parser){
    if (!(parser = av_parser_init(id))) {
      printf("parser not found\n");
      return false;
    }
  }
  if(!videoCodecCtx){
    if(!(videoCodecCtx = avcodec_alloc_context3(videoCodec))){
      printf("Cannot find valid decode codec context.\n");
      return false;
    }
  }

  if(!openCodec){
    if (avcodec_open2(videoCodecCtx, videoCodec, NULL) < 0) {
      printf("Cannot open codec.\n");
      return false;
    }
    openCodec = true;
  }

  return true;
}

void FFmpegDecoder::parsePkt(uint8_t *data, int len)
{
  int eof = !len;
  while (len > 0 || eof) {
    int ret = av_parser_parse2(parser, videoCodecCtx, &pkt->data, &pkt->size,
                               data, 4096, AV_NOPTS_VALUE, AV_NOPTS_VALUE, 0);
    if (ret < 0) {
      fprintf(stderr, "Error while parsing\n");
      continue;
    }

    data += ret;
    len -= ret;

    if (pkt->size) {
      decode();
    }
  }
}

void FFmpegDecoder::decode()
{
  int ret = avcodec_send_packet(videoCodecCtx, pkt);
  if (ret >= 0) {
    while ((ret = avcodec_receive_frame(videoCodecCtx, yuvFrame)) >= 0) {
      if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
        break;
      else if (ret < 0) {
        fprintf(stderr, "Error during decoding\n");
        break;
      }
      if(!img_ctx){
        img_ctx = sws_getContext(videoCodecCtx->width,
                                 videoCodecCtx->height,
                                 videoCodecCtx->pix_fmt,
                                 videoCodecCtx->width,
                                 videoCodecCtx->height,
                                 AV_PIX_FMT_RGB32,
                                 SWS_BICUBIC, NULL, NULL, NULL);
      }
      if(!out_buffer){
        int numBytes = av_image_get_buffer_size(AV_PIX_FMT_RGB32,videoCodecCtx->width,videoCodecCtx->height,1);
        out_buffer = (unsigned char *)av_malloc(numBytes*sizeof(unsigned char));

        int res = av_image_fill_arrays(
              rgbFrame->data,rgbFrame->linesize,
              out_buffer,AV_PIX_FMT_RGB32,
              videoCodecCtx->width,videoCodecCtx->height,1);
        if(res<0){
          qDebug()<<"Fill arrays failed.";
          break;
        }
      }

      sws_scale(img_ctx,
                yuvFrame->data, yuvFrame->linesize,
                0, videoCodecCtx->height,
                rgbFrame->data, rgbFrame->linesize);

      QImage img(out_buffer,
                 videoCodecCtx->width, videoCodecCtx->height,
                 QImage::Format_RGB32);
      setImage(img);
      av_packet_unref(pkt);
    }
  }
  av_packet_unref(pkt);
}

QImage FFmpegDecoder::getImage() const
{
  return image;
}

void FFmpegDecoder::parsePkt(AVPacket& packet)
{
  pkt->data = packet.data;
  pkt->size = packet.size;
  decode();
}

void FFmpegDecoder::clear()
{
  avcodec_flush_buffers(videoCodecCtx);
}

void FFmpegDecoder::setImage(const QImage &newImage)
{
  QMutexLocker locker(&mutex);
  image = newImage;
}