// Copyright (c) 2012 The Foundry Visionmongers Ltd.  All Rights Reserved.

#include "Build/fnBuild.h"
#include "DDImage/DDString.h"
#include "DDImage/Writer.h"
#include "DDImage/Row.h"
#include "DDImage/Knobs.h"

#define INT64_C(c) (c ## LL)
#define UINT64_C(c) (c ## ULL)

extern "C" {
#include <errno.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include "libavutil/imgutils.h"
#include "libavformat/avio.h"
#include <libavutil/opt.h>
#include <libswscale/swscale.h>
#include <libavutil/avutil.h>
#include <libavutil/error.h>
}

using namespace DD::Image;

class ffmpegWriter : public Writer
{
private:
  enum WriterError { SUCCESS = 0, IGNORE_FINISH, CLEANUP };
  
public:
  explicit ffmpegWriter(Write* iop);
  ~ffmpegWriter();
  
  virtual bool movie() const { return true; }
  
  void execute();
  void finish();
  void knobs(Knob_Callback f);
  static const Writer::Description d;
  
private:
  void freeFormat();
  bool isCodecBlacklisted(const char* name);
  
private:
  AVCodecContext*   codecContext_;
  AVFormatContext*  formatContext_;
  AVStream* stream_;
  std::vector<std::string> formatsLongNames_;
  std::vector<const char*> formatsShortNames_;
  std::vector<const char*> codecsLongNames_;
  std::vector<const char*> codecsShortNames_;
  
  WriterError error_;
  // knobs variables
  float fps_;
  int format_;
  int codec_;
  int bitrate_;
  int bitrateTolerance_;
  int gopSize_;
  int bFrames_;
  int mbDecision_;
};

static const char* codecBlacklist[] =
{
  "r10k",
  "ogg",
  "mpjpeg",
  "asf",
  "asf_stream",
  "h261",
  "h263",
  "rcv",
  "yuv4mpegpipe",
  "prores",
  "dnxhd",
  "swf",
  "dv",
  "ipod",
  "psp",
  "image2",
  "3g2",
  "3gp",
  "RoQ",
  
  #ifdef FN_OS_WINDOWS
  
  "dirac",
  "ffm",
  
  #elif FN_OS_MAC
  
  "mjpeg",
  "vc1",
  
  #elif FN_OS_LINUX
  
  "mjpeg",
  "vc1",  
  
  #endif
  
  NULL
  
};

ffmpegWriter::ffmpegWriter(Write* iop)
  : Writer(iop)
  , codecContext_(NULL)
  , formatContext_(NULL)
  , stream_(0)
  , error_(IGNORE_FINISH)
  , fps_(DD::Image::root_real_fps())
  , format_(0)
  , codec_(0)
  , bitrate_(400000)
  , bitrateTolerance_(4000 * 10000)
  , gopSize_(12)
  , bFrames_(0)
  , mbDecision_(FF_MB_DECISION_SIMPLE)
{
  av_log_set_level(AV_LOG_WARNING);
  av_register_all();
  
  formatsLongNames_.push_back("default");
  formatsShortNames_.push_back("default");
  AVOutputFormat* fmt = av_oformat_next(NULL);
  while (fmt) {
    
    if (fmt->video_codec != CODEC_ID_NONE && !isCodecBlacklisted( fmt->name ) ) {
      if (fmt->long_name) {
        formatsLongNames_.push_back(std::string(fmt->long_name) + std::string(" (") + std::string(fmt->name) + std::string(")"));
        formatsShortNames_.push_back(fmt->name);
      }
    }
    fmt = av_oformat_next(fmt);
  }
  formatsShortNames_.push_back(0);
  
  codecsLongNames_.push_back("default");
  codecsShortNames_.push_back("default");
  AVCodec* c = av_codec_next(NULL);
  while (c) {
    if (c->type == AVMEDIA_TYPE_VIDEO && c->encode && !isCodecBlacklisted( c->name ) ) {
      if (c->long_name) {
        codecsLongNames_.push_back(c->long_name);
        codecsShortNames_.push_back(c->name);
      }
    }
    c = av_codec_next(c);
  }
  codecsLongNames_.push_back(0);
  codecsShortNames_.push_back(0);
}

bool ffmpegWriter::isCodecBlacklisted(const char* name)
{
  const char** iterator = codecBlacklist;
  
  while( *iterator != NULL ) {
    if (strncmp(name, *iterator, strlen(*iterator)) == 0)  {
      return true;
    }
    
    ++iterator;
  }
  return false;
}

ffmpegWriter::~ffmpegWriter()
{
  av_free(codecContext_);
}

void ffmpegWriter::execute()
{
  error_ = IGNORE_FINISH;
  
  AVOutputFormat* fmt = 0;
  if (!format_) {
    fmt = av_guess_format(NULL, filename(), NULL);
    if (!fmt) {
      iop->critical("could not deduce output format from file extension");
      return;
    }
  }
  else {
    fmt = av_guess_format(formatsShortNames_[format_], NULL, NULL);
    if (!fmt) {
      iop->critical("could not deduce output format");
      return;
    }
  }
  
  if (!formatContext_)
    avformat_alloc_output_context2(&formatContext_, fmt, NULL, filename());

  snprintf(formatContext_->filename, sizeof(formatContext_->filename), "%s", filename());

  CodecID codecId = fmt->video_codec;
  if (codec_) {
    AVCodec* userCodec = avcodec_find_encoder_by_name(codecsShortNames_[codec_]);
    if (userCodec) {
      codecId = userCodec->id;
    }
  }
  
  AVCodec* videoCodec = avcodec_find_encoder(codecId);
  if (!videoCodec) {
    iop->critical("unable to find codec");
    freeFormat();
    return;
  }
  
  PixelFormat pixFMT = PIX_FMT_YUV420P;
  
  if (videoCodec->pix_fmts != NULL) {
    pixFMT = *videoCodec->pix_fmts;
  }
  else {
    if (strcmp(fmt->name, "gif") == 0) {
      pixFMT = PIX_FMT_RGB24;
    }
  }
  
  bool isCodecSupportedInContainer = (avformat_query_codec(fmt, codecId, FF_COMPLIANCE_NORMAL) == 1);
  // mov seems to be able to cope with anything, which the above function doesn't seem to think is the case (even with FF_COMPLIANCE_EXPERIMENTAL)
  // and it doesn't return -1 for in this case, so we'll special-case this situation to allow this
  isCodecSupportedInContainer |= (strcmp(formatContext_->oformat->name, "mov") == 0);
  
  if (!isCodecSupportedInContainer) {
    iop->critical("the selected codec is not supported in this container.");
    freeFormat();
    return;
  }
  
  if (isCodecBlacklisted(videoCodec->name)) {
    iop->critical("unsupported codec");
    freeFormat();
    return;
  }
    
  if (!stream_) {
    stream_ = avformat_new_stream(formatContext_, NULL);
    if (!stream_) {
      iop->critical("out of memory");
      return;
    }

    codecContext_ = stream_->codec;
    
    // this seems to be needed for certain codecs, as otherwise they don't have relevant options set
    avcodec_get_context_defaults3(codecContext_, videoCodec);
    
    codecContext_->pix_fmt = pixFMT;   // this is set to the first element of FMT a choice could be added
    
    codecContext_->bit_rate = bitrate_;
    codecContext_->bit_rate_tolerance = bitrateTolerance_;
    codecContext_->width = width();
    codecContext_->height = height();
    
    // Bug 23953
    // ffmpeg does a horrible job of converting floats to AVRationals
    // It adds 0.5 randomly and does some other stuff
    // To work around that, we just multiply the fps by what I think is a reasonable number to make it an int
    // and use the reasonable number as the numerator for the timebase.
    // Timebase is not the frame rate; it's the inverse of the framerate
    // So instead of doing 1/fps, we just set the numerator and denominator of the timebase directly.
    // The upshot is that this allows ffmpeg to properly do framerates of 23.78 (or 23.796, which is what the user really wants when they put that in).
    //
    // The code was this:
    //stream_->codec->time_base = av_d2q(1.0 / fps_, 100);
    const float CONVERSION_FACTOR = 1000.0f;
    codecContext_->time_base.num = (int) CONVERSION_FACTOR;
    codecContext_->time_base.den = (int) (fps_ * CONVERSION_FACTOR);
    //
    
    codecContext_->gop_size = gopSize_;
    // NOTE: in new ffmpeg, bframes don't seem to work correctly - ffmpeg crashes...
    if (bFrames_) {
      codecContext_->max_b_frames = bFrames_;
      codecContext_->b_frame_strategy = 0;
      codecContext_->b_quant_factor = 2.0f;
    }
      
    codecContext_->mb_decision = mbDecision_;
    
    if (!strcmp(formatContext_->oformat->name, "mp4") || !strcmp(formatContext_->oformat->name, "mov") || !strcmp(formatContext_->oformat->name, "3gp"))
      codecContext_->flags |= CODEC_FLAG_GLOBAL_HEADER;
    
    if (formatContext_->oformat->flags & AVFMT_GLOBALHEADER)
      codecContext_->flags |= CODEC_FLAG_GLOBAL_HEADER;
      
    if (avcodec_open2(codecContext_, videoCodec, NULL) < 0) {
      iop->critical("unable to open codec");
      freeFormat();
      return;
    }
                  
    if (!(fmt->flags & AVFMT_NOFILE)) {
      if (avio_open(&formatContext_->pb, filename(), AVIO_FLAG_WRITE) < 0) {
        iop->critical("unable to open file");
        freeFormat();
        return;
      }
    }
    
    avformat_write_header(formatContext_, NULL);
  }
  
  error_ = CLEANUP;
  
  AVPicture picture;
  int picSize = avpicture_get_size(PIX_FMT_RGB24, width(), height());
  // allocate a buffer for the picture's image...
  uint8_t* buffer = (uint8_t*)av_malloc(picSize);
  // blank the values - this initialises stuff and seems to be needed
  avpicture_fill(&picture, buffer, PIX_FMT_RGB24, width(), height());
    
  Row row(0, width());
  input0().validate();
  input0().request(0, 0, width(), height(), Mask_RGB, 1);
 
  for (int y = 0; y < height(); ++y) {
    get(y, 0, width(), Mask_RGB, row);
    if (iop->aborted()) {
      av_free(buffer);
      return;
    }
    
    for (Channel z = Chan_Red; z <= Chan_Blue; incr(z)) {
      const float* from = row[z];
      to_byte(z - 1, picture.data[0] + (height() - y - 1) * picture.linesize[0] + z - 1, from, NULL, width(), 3);
    }
  }
  
  // now allocate an image frame for the image in the output codec's format...
  AVFrame* output = avcodec_alloc_frame();
  picSize = avpicture_get_size(pixFMT, width(), height());
  uint8_t* outBuffer = (uint8_t*)av_malloc(picSize);
  
  av_image_alloc(output->data, output->linesize, width(), height(), pixFMT, 1);
  
  SwsContext* convertCtx = sws_getContext(width(), height(), PIX_FMT_RGB24, width(), height(),
                                                pixFMT, SWS_BICUBIC, NULL, NULL, NULL);
  
  int sliceHeight = sws_scale(convertCtx, picture.data, picture.linesize, 0, height(), output->data, output->linesize);
  assert(sliceHeight > 0);
  
  int ret = 0;
  if ((formatContext_->oformat->flags & AVFMT_RAWPICTURE) != 0) {
    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.flags |= AV_PKT_FLAG_KEY;
    pkt.stream_index = stream_->index;
    pkt.data = (uint8_t*)output;
    pkt.size = sizeof(AVPicture);
    ret = av_interleaved_write_frame(formatContext_, &pkt);
  }
  else {
    uint8_t* outbuf = (uint8_t*)av_malloc(picSize);
    assert(outbuf != NULL);
    ret = avcodec_encode_video(codecContext_, outbuf, picSize, output);
    if (ret > 0) {
      AVPacket pkt;
      av_init_packet(&pkt);
      if (codecContext_->coded_frame && static_cast<unsigned long>(codecContext_->coded_frame->pts) != AV_NOPTS_VALUE)
        pkt.pts = av_rescale_q(codecContext_->coded_frame->pts, codecContext_->time_base, stream_->time_base);
      if (codecContext_->coded_frame && codecContext_->coded_frame->key_frame)
        pkt.flags |= AV_PKT_FLAG_KEY;
      
      pkt.stream_index = stream_->index;
      pkt.data = outbuf;
      pkt.size = ret;
      
      ret = av_interleaved_write_frame(formatContext_, &pkt);
    }
    else {
      // we've got an error
      char szError[1024];
      av_strerror(ret, szError, 1024);
      iop->error(szError);
    }
    
    av_free(outbuf);
  }
  
  av_free(outBuffer);
  av_free(buffer);
  av_free(output);
  
  if (ret) {
    iop->critical("error writing frame to file");
    return;
  }
  
  error_ = SUCCESS;
}

void ffmpegWriter::finish()
{
  if (error_ == IGNORE_FINISH)
    return;
  av_write_trailer(formatContext_);
  avcodec_close(codecContext_);
  if (!(formatContext_->oformat->flags & AVFMT_NOFILE))
    avio_close(formatContext_->pb);
  freeFormat();
}

void ffmpegWriter::knobs(Knob_Callback f)
{
  static std::vector<const char*> formatsAliases;
  
  formatsAliases.resize(formatsLongNames_.size());
  for (int i = 0; i < static_cast<int>(formatsLongNames_.size()); ++i)
    formatsAliases[i] = formatsLongNames_[i].c_str();
  formatsAliases.push_back(0);
  
  Enumeration_knob(f, &format_, &formatsAliases[0], "format");
  Float_knob(f, &fps_, IRange(0.0, 100.0f), "fps");
  
  BeginClosedGroup(f, "Advanced");
  
  Enumeration_knob(f, &codec_, &codecsLongNames_[0], "codec");
  Int_knob(f, &bitrate_, IRange(0.0, 400000), "bitrate");
  SetFlags(f, Knob::SLIDER | Knob::LOG_SLIDER);
  Int_knob(f, &bitrateTolerance_, IRange(0, 4000 * 10000), "bitrateTol", "bitrate tolerance");
  SetFlags(f, Knob::SLIDER | Knob::LOG_SLIDER);
  Int_knob(f, &gopSize_, IRange(0, 30), "gopSize", "GOP size");
  SetFlags(f, Knob::SLIDER | Knob::LOG_SLIDER);
  Int_knob(f, &bFrames_, IRange(0, 30), "bFrames", "B Frames");
  SetFlags(f, Knob::SLIDER | Knob::LOG_SLIDER);
  
  static const char* mbDecisionTypes[] = {
    "FF_MB_DECISION_SIMPLE", "FF_MB_DECISION_BITS", "FF_MB_DECISION_RD", 0
  };
  Enumeration_knob(f, &mbDecision_, mbDecisionTypes, "mbDecision", "macro block decision mode");
  
  EndGroup(f);
}

void ffmpegWriter::freeFormat()
{
  for (int i = 0; i < static_cast<int>(formatContext_->nb_streams); ++i)
    av_freep(&formatContext_->streams[i]);
  av_free(formatContext_);
  formatContext_ = NULL;
  stream_ = NULL;
}

static Writer* build(Write* iop)
{
  return new ffmpegWriter(iop);
}

const Writer::Description ffmpegWriter::d("ffmpeg\0mov\0avi\0mp4\0m4v\0", build);
