/* This file is an image processing operation for GEGL
 *
 * GEGL 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.
 *
 * GEGL 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 GEGL; if not, see <http://www.gnu.org/licenses/>.
 *
 * Copyright 2003,2004,2007 Øyvind Kolås <pippin@gimp.org>
 */

#include "config.h"

#include <stdlib.h>

#include <glib/gi18n-lib.h>

#include <libswscale/swscale.h>


#ifdef GEGL_CHANT_PROPERTIES

gegl_chant_string (path, _("File"), "/tmp/fnord.mp4", _("Target path and filename, use '-' for stdout."))
gegl_chant_double (bitrate, _("Bitrate"), 0.0, 100000000.0, 800000.0, _("target bitrate"))
gegl_chant_double (fps, _("FPS"), 0.0, 100.0, 25, _("frames per second"))

#else

#define GEGL_CHANT_TYPE_SINK
#define GEGL_CHANT_C_FILE       "ff-save.c"

#include "gegl-chant.h"

#ifdef HAVE_LIBAVFORMAT_AVFORMAT_H
#include <libavformat/avformat.h>
#else
#include <avformat.h>
#endif

typedef struct
{
  gdouble    frame;
  gdouble    frames;
  gdouble    width;
  gdouble    height;
  GeglBuffer *input;

  AVOutputFormat *fmt;
  AVFormatContext *oc;
  AVStream *video_st;

  AVFrame  *picture, *tmp_picture;
  uint8_t  *video_outbuf;
  int       frame_count, video_outbuf_size;

    /** the rest is for audio handling within oxide, note that the interface
     * used passes all used functions in the oxide api through the reg_sym api
     * of gggl, this means that the ops should be usable by other applications
     * using gggl directly,. without needing to link with the oxide library
     */

  AVStream *audio_st;

  void     *oxide_audio_instance;
  /*< non NULL audio_query,. means audio present */

  int32_t (*oxide_audio_query) (void *audio_instance,
                                uint32_t * sample_rate,
                                uint32_t * bits,
                                uint32_t * channels,
                                uint32_t * fragment_samples,
                                uint32_t * fragment_size);

  /* get audio samples for the current video frame, this should provide all
   * audiosamples associated with the frame, frame centering on audio stream is
   * undefined (FIXME:<<)
   */

  int32_t (*oxide_audio_get_fragment) (void *audio_instance, uint8_t * buf);

  uint32_t  sample_rate;
  uint32_t  bits;
  uint32_t  channels;
  uint32_t  fragment_samples;
  uint32_t  fragment_size;

  uint8_t  *fragment;

  int       buffer_size;
  int       buffer_read_pos;
  int       buffer_write_pos;
  uint8_t  *buffer;   /* optimal buffer size should be calculated,. preferably
                         run time,. no magic numbers needed for the filewrite
                         case, perhaps a tiny margin for ntsc since it has a
                         strange framerate */

  int       audio_outbuf_size;
  int       audio_input_frame_size;
  int16_t  *samples;
  uint8_t  *audio_outbuf;
} Priv;

#define DISABLE_AUDIO

static void
init (GeglChantO *o)
{
  static gint inited = 0; /*< this is actually meant to be static, only to be done once */
  Priv       *p = (Priv*)o->chant_data;

  if (p == NULL)
    {
      p = g_new0 (Priv, 1);
      o->chant_data = (void*) p;
    }

  if (!inited)
    {
      av_register_all ();
      avcodec_register_all ();
      inited = 1;
    }

#ifndef DISABLE_AUDIO
  p->oxide_audio_instance = gggl_op_sym (op, "oxide_audio_instance");
  p->oxide_audio_query = gggl_op_sym (op, "oxide_audio_query()");
  p->oxide_audio_get_fragment =
    gggl_op_sym (op, "oxide_audio_get_fragment()");

  if (p->oxide_audio_instance && p->oxide_audio_query)
    {
      p->oxide_audio_query (p->oxide_audio_instance,
                            &p->sample_rate,
                            &p->bits,
                            &p->channels,
                            &p->fragment_samples, &p->fragment_size);

      /* FIXME: for now, the buffer is set to a size double that of a oxide
       * provided fragment,. should be enough no matter how things are handled,
       * but it should also be more than needed,. find out exact amount needed later
       */

      if (!p->buffer)
        {
          int size =
            (p->sample_rate / p->fps) * p->channels * (p->bits / 8) * 2;
          buffer_open (op, size);
        }

      if (!p->fragment)
        p->fragment = gggl_op_calloc (op, 1, p->fragment_size);
    }
#endif
}

static void close_video       (Priv            *p,
                               AVFormatContext *oc,
                               AVStream        *st);
void        close_audio       (Priv            *p,
                               AVFormatContext *oc,
                               AVStream        *st);
static int  tfile             (GeglChantO      *self);
static void write_video_frame (GeglChantO      *self,
                               AVFormatContext *oc,
                               AVStream        *st);
static void write_audio_frame (GeglChantO      *self,
                               AVFormatContext *oc,
                               AVStream        *st);




#define STREAM_FRAME_RATE 25    /* 25 images/s */

static int
buffer_used (Priv * p)
{
  int       ret;
  if (!p || !p->buffer_size || !p->buffer)
    return -1;

  if (p->buffer_write_pos == p->buffer_read_pos)
    return 0;
  if (p->buffer_write_pos > p->buffer_read_pos)
    {
      ret = p->buffer_write_pos - p->buffer_read_pos;
    }
  else
    {
      ret = p->buffer_size - (p->buffer_read_pos - p->buffer_write_pos);
    }
  return ret;
}

static int
buffer_unused (Priv * p)
{
  int       ret;
  ret = p->buffer_size - buffer_used (p) - 1;   /* 1 byte for indicating full */
  return ret;
}

#ifndef DISABLE_AUDIO
static void
buffer_flush (Priv * p)
{
  if (!p->buffer)
    return;
  p->buffer_write_pos = 0;
  p->buffer_read_pos = 0;
}

static void
buffer_open (GeglChantOperation *op, int size)
{
  Priv     *p = (Priv*)op->priv;

  if (p->buffer)
    {
      fprintf (stderr, "double init of buffer, eek!\n");
    }

  p->buffer_size = size;
  p->buffer = g_malloc (size);
  buffer_flush (p);
}

static void
buffer_close (GeglChantOperation *op)
{
  Priv     *p = (Priv*)op->priv;

  if (!p->buffer)
    return;
  g_free (p->buffer);
  p->buffer = NULL;
}

#endif

static int
buffer_write (Priv * p, uint8_t * source, int count)
{
  int       first_segment_size = 0;
  int       second_segment_size = 0;

  /* check if we have room for the data in the buffer */
  if (!p->buffer || buffer_unused (p) < count + 1)
    {
      fprintf (stderr, "ff_save audio buffer full!! shouldn't happen\n");
      return -1;
    }

  /* calculate size of segments to write */
  first_segment_size = p->buffer_size - p->buffer_write_pos;

  if (p->buffer_read_pos > p->buffer_write_pos)
    first_segment_size = p->buffer_read_pos - p->buffer_write_pos;

  if (first_segment_size >= count)
    {
      first_segment_size = count;
    }
  else
    {
      second_segment_size = count - first_segment_size;
    }

  memcpy (p->buffer + p->buffer_write_pos, source, first_segment_size);
  p->buffer_write_pos += first_segment_size;

  if (p->buffer_write_pos == p->buffer_size)
    p->buffer_write_pos = 0;

  if (second_segment_size)
    {
      memcpy (p->buffer + p->buffer_write_pos, source + first_segment_size,
              second_segment_size);
      p->buffer_write_pos = second_segment_size;
    }

  return first_segment_size + second_segment_size;
}

#if 0
. = unused byte
# = used byte
  all sizes are specified in contiguous bytes of memory
  __write_pos / __read_pos
  //
  |................................... | empty buffer
  \ _________________________________ /
  size_ / __read_pos __write_pos / /|........
#################..........|
  \ ______________ / used ()_ / __write_pos __read_pos / /|
############..................#####|
  \______________ / unused ()_ / ___write_pos / __read_pos / /|
#################.#################|    full buffer
#endif
     static int
     buffer_read (Priv * p, uint8_t * dest, int count)
{
  int       first_segment_size = 0;
  int       second_segment_size = 0;

  /* check if we have enough data to fulfil request */
  if (!p->buffer || buffer_used (p) < count)
    {
      fprintf (stderr, "ff_save audio buffer doesn't have enough data\n");
      return -1;
    }

  /* calculate size of segments to write */
  first_segment_size = p->buffer_size - p->buffer_read_pos;

  if (p->buffer_write_pos > p->buffer_read_pos)
    first_segment_size = p->buffer_write_pos - p->buffer_read_pos;

  if (first_segment_size >= count)
    {
      first_segment_size = count;
    }
  else
    {
      second_segment_size = count - first_segment_size;
    }

  memcpy (dest, p->buffer + p->buffer_read_pos, first_segment_size);
  p->buffer_read_pos += first_segment_size;

  if (p->buffer_read_pos == p->buffer_size)
    p->buffer_read_pos = 0;

  if (second_segment_size)
    {
      memcpy (dest + first_segment_size, p->buffer + p->buffer_read_pos,
              second_segment_size);
      p->buffer_read_pos = second_segment_size;
    }

  return first_segment_size + second_segment_size;
}

#ifndef DISABLE_AUDIO
/* add an audio output stream */
static AVStream *
add_audio_stream (GeglChantOperation *op, AVFormatContext * oc, int codec_id)
{
  Priv     *p = (Priv*)op->priv;
  AVCodecContext *c;
  AVStream *st;

  p = NULL;
  st = av_new_stream (oc, 1);
  if (!st)
    {
      fprintf (stderr, "Could not alloc stream\n");
      exit (1);
    }

  c = st->codec;
  c->codec_id = codec_id;
  c->codec_type = CODEC_TYPE_AUDIO;

  c->bit_rate = 64000;
  c->sample_rate = 44100;
  c->channels = 2;
  return st;
}
#endif

static void
open_audio (Priv * p, AVFormatContext * oc, AVStream * st)
{
  AVCodecContext *c;
  AVCodec  *codec;

  c = st->codec;

  /* find the audio encoder */
  codec = avcodec_find_encoder (c->codec_id);
  if (!codec)
    {
      fprintf (stderr, "codec not found\n");
      exit (1);
    }

  /* open it */
  if (avcodec_open2 (c, codec, NULL) < 0)
    {
      fprintf (stderr, "could not open codec\n");
      exit (1);
    }

  p->audio_outbuf_size = 10000;
  p->audio_outbuf = malloc (p->audio_outbuf_size);

  /* ugly hack for PCM codecs (will be removed ASAP with new PCM
     support to compute the input frame size in samples */
  if (c->frame_size <= 1)
    {
      fprintf (stderr, "eeko\n");
      p->audio_input_frame_size = p->audio_outbuf_size / c->channels;
      switch (st->codec->codec_id)
        {
        case CODEC_ID_PCM_S16LE:
        case CODEC_ID_PCM_S16BE:
        case CODEC_ID_PCM_U16LE:
        case CODEC_ID_PCM_U16BE:
          p->audio_input_frame_size >>= 1;
          break;
        default:
          break;
        }
    }
  else
    {
      p->audio_input_frame_size = c->frame_size;
    }
  /*audio_input_frame_size = 44100/25;*/
  p->samples = malloc (p->audio_input_frame_size * 2 * c->channels);
}

void
write_audio_frame (GeglChantO *op, AVFormatContext * oc, AVStream * st)
{
  Priv *p = (Priv*)op->chant_data;

  AVCodecContext *c;
  AVPacket  pkt;
  av_init_packet (&pkt);

  c = st->codec;

  /*fprintf (stderr, "going to grab %i\n", p->fragment_size);*/
  if (p->oxide_audio_get_fragment (p->oxide_audio_instance,
                                   p->fragment) == (signed) p->fragment_size)
    {
      buffer_write (p, p->fragment, p->fragment_size);
    }

  while (buffer_used (p) >= p->audio_input_frame_size * 2 * c->channels)
    {
      buffer_read (p, (uint8_t *) p->samples,
                   p->audio_input_frame_size * 2 * c->channels);

      pkt.size = avcodec_encode_audio (c, p->audio_outbuf,
                                       p->audio_outbuf_size, p->samples);

      pkt.pts = c->coded_frame->pts;
      pkt.flags |= AV_PKT_FLAG_KEY;
      pkt.stream_index = st->index;
      pkt.data = p->audio_outbuf;

      if (av_write_frame (oc, &pkt) != 0)
        {
          fprintf (stderr, "Error while writing audio frame\n");
          exit (1);
        }
    }
}

/*p->audio_get_frame (samples, audio_input_frame_size, c->channels);*/

void
close_audio (Priv * p, AVFormatContext * oc, AVStream * st)
{
  avcodec_close (st->codec);

  av_free (p->samples);
  av_free (p->audio_outbuf);
}

/* add a video output stream */
static AVStream *
add_video_stream (GeglChantO *op, AVFormatContext * oc, int codec_id)
{
  Priv *p = (Priv*)op->chant_data;

  AVCodecContext *c;
  AVStream *st;

  st = av_new_stream (oc, 0);
  if (!st)
    {
      fprintf (stderr, "Could not alloc stream %p %p %i\n", op, oc, codec_id);
      exit (1);
    }

  c = st->codec;
  c->codec_id = codec_id;
  c->codec_type = AVMEDIA_TYPE_VIDEO;

  /* put sample propeters */
  c->bit_rate = op->bitrate;
  /* resolution must be a multiple of two */
  c->width = p->width;
  c->height = p->height;
  /* frames per second */
  /*c->frame_rate = op->fps;
  c->frame_rate_base = 1;*/

#if LIBAVCODEC_BUILD >= 4754
  c->time_base=(AVRational){1, op->fps};
    #else
        c->frame_rate=op->fps;
        c->frame_rate_base=1;
    #endif
     c->pix_fmt = PIX_FMT_YUV420P;


  c->gop_size = 12;             /* emit one intra frame every twelve frames at most */
  if (c->codec_id == CODEC_ID_MPEG2VIDEO)
    {
      /* just for testing, we also add B frames */
      /*c->max_b_frames = 2;*/
    }
/*    if (!strcmp (oc->oformat->name, "mp4") ||
          !strcmp (oc->oformat->name, "3gp"))
    c->flags |= CODEC_FLAG_GLOBAL_HEADER;
    */
  return st;
}


static AVFrame *
alloc_picture (int pix_fmt, int width, int height)
{
  AVFrame  *picture;
  uint8_t  *picture_buf;
  int       size;

  picture = avcodec_alloc_frame ();
  if (!picture)
    return NULL;
  size = avpicture_get_size (pix_fmt, width, height);
  picture_buf = malloc (size);
  if (!picture_buf)
    {
      av_free (picture);
      return NULL;
    }
  avpicture_fill ((AVPicture *) picture, picture_buf, pix_fmt, width, height);
  return picture;
}

static void
open_video (Priv * p, AVFormatContext * oc, AVStream * st)
{
  AVCodec  *codec;
  AVCodecContext *c;

  c = st->codec;

  /* find the video encoder */
  codec = avcodec_find_encoder (c->codec_id);
  if (!codec)
    {
      fprintf (stderr, "codec not found\n");
      exit (1);
    }

  /* open the codec */
  if (avcodec_open (c, codec) < 0)
    {
      fprintf (stderr, "could not open codec\n");
      exit (1);
    }

  p->video_outbuf = NULL;
  if (!(oc->oformat->flags & AVFMT_RAWPICTURE))
    {
      /* allocate output buffer */
      /* XXX: API change will be done */
      p->video_outbuf_size = 200000;
      p->video_outbuf = malloc (p->video_outbuf_size);
    }

  /* allocate the encoded raw picture */
  p->picture = alloc_picture (c->pix_fmt, c->width, c->height);
  if (!p->picture)
    {
      fprintf (stderr, "Could not allocate picture\n");
      exit (1);
    }

  /* if the output format is not YUV420P, then a temporary YUV420P
     picture is needed too. It is then converted to the required
     output format */
  p->tmp_picture = NULL;
  if (c->pix_fmt != PIX_FMT_RGB24)
    {
      p->tmp_picture = alloc_picture (PIX_FMT_RGB24, c->width, c->height);
      if (!p->tmp_picture)
        {
          fprintf (stderr, "Could not allocate temporary picture\n");
          exit (1);
        }
    }
}

static void
close_video (Priv * p, AVFormatContext * oc, AVStream * st)
{
  avcodec_close (st->codec);
  av_free (p->picture->data[0]);
  av_free (p->picture);
  if (p->tmp_picture)
    {
      av_free (p->tmp_picture->data[0]);
      av_free (p->tmp_picture);
    }
  av_free (p->video_outbuf);
}

#include "string.h"

/* prepare a dummy image */
static void
fill_yuv_image (GeglChantO *op,
                AVFrame *pict, int frame_index, int width, int height)
{
  Priv     *p = (Priv*)op->chant_data;
  /*memcpy (pict->data[0],

   op->input_pad[0]->data,

          op->input_pad[0]->width * op->input_pad[0]->height * 3);*/
  GeglRectangle rect={0,0,width,height};
  gegl_buffer_get (p->input, &rect, 1.0, babl_format ("R'G'B' u8"), pict->data[0], GEGL_AUTO_ROWSTRIDE, GEGL_ABYSS_NONE);
}

static void
write_video_frame (GeglChantO *op,
                   AVFormatContext *oc, AVStream *st)
{
  Priv     *p = (Priv*)op->chant_data;
  int       out_size, ret;
  AVCodecContext *c;
  AVFrame  *picture_ptr;

  c = st->codec;

  if (c->pix_fmt != PIX_FMT_RGB24)
    {
      struct SwsContext *img_convert_ctx;

      /* as we only generate a RGB24 picture, we must convert it
         to the codec pixel format if needed */
      fill_yuv_image (op, p->tmp_picture, p->frame_count, c->width,
                      c->height);

      img_convert_ctx = sws_getContext(c->width, c->height, c->pix_fmt,
                                       c->width, c->height, PIX_FMT_RGB24,
                                       SWS_BICUBIC, NULL, NULL, NULL);

      if (img_convert_ctx == NULL)
        {
          fprintf(stderr, "ff_save: Cannot initialize conversion context.");
        }
      else
        {
          sws_scale(img_convert_ctx,
                    p->tmp_picture->data,
                    p->tmp_picture->linesize,
                    0,
                    c->height,
                    p->picture->data,
                    p->picture->linesize);
        }
    }
  else
    {
      fill_yuv_image (op, p->picture, p->frame_count, c->width, c->height);
    }
  picture_ptr = p->picture;

  if (oc->oformat->flags & AVFMT_RAWPICTURE)
    {
      /* raw video case. The API will change slightly in the near
         future for that */
      AVPacket  pkt;
      av_init_packet (&pkt);

      pkt.flags |= AV_PKT_FLAG_KEY;
      pkt.stream_index = st->index;
      pkt.data = (uint8_t *) picture_ptr;
      pkt.size = sizeof (AVPicture);

      ret = av_write_frame (oc, &pkt);
    }
  else
    {
      /* encode the image */
      out_size =
        avcodec_encode_video (c,
                              p->video_outbuf,
                              p->video_outbuf_size, picture_ptr);

      /* if zero size, it means the image was buffered */
      if (out_size != 0)
        {
          AVPacket  pkt;
          av_init_packet (&pkt);

          pkt.pts = c->coded_frame->pts;
          if (c->coded_frame->key_frame)
            pkt.flags |= AV_PKT_FLAG_KEY;
          pkt.stream_index = st->index;
          pkt.data = p->video_outbuf;
          pkt.size = out_size;

          /* write the compressed frame in the media file */
          ret = av_write_frame (oc, &pkt);
        }
      else
        {
          ret = 0;
        }
    }
  if (ret != 0)
    {
      fprintf (stderr, "Error while writing video frame\n");
      exit (1);
    }
  p->frame_count++;
}

static int
tfile (GeglChantO *self)
{
  Priv *p = (Priv*)self->chant_data;

  p->fmt = av_guess_format (NULL, self->path, NULL);
  if (!p->fmt)
    {
      fprintf (stderr,
               "ff_save couldn't deduce outputformat from file extension: using MPEG.\n%s",
               "");
      p->fmt = av_guess_format ("mpeg", NULL, NULL);
    }
  p->oc = avformat_alloc_context ();/*g_malloc (sizeof (AVFormatContext));*/
  if (!p->oc)
    {
      fprintf (stderr, "memory error\n%s", "");
      return -1;
    }

  p->oc->oformat = p->fmt;

  snprintf (p->oc->filename, sizeof (p->oc->filename), "%s", self->path);

  p->video_st = NULL;
  p->audio_st = NULL;

  if (p->fmt->video_codec != CODEC_ID_NONE)
    {
      p->video_st = add_video_stream (self, p->oc, p->fmt->video_codec);
    }
  if (p->oxide_audio_query && p->fmt->audio_codec != CODEC_ID_NONE)
    {
     /*XXX: FOO p->audio_st = add_audio_stream (op, p->oc, p->fmt->audio_codec);*/
    }

  if (av_set_parameters (p->oc, NULL) < 0)
    {
      fprintf (stderr, "Invalid output format propeters\n%s", "");
      return -1;
    }

  dump_format (p->oc, 0, self->path, 1);

  if (p->video_st)
    open_video (p, p->oc, p->video_st);
  if (p->audio_st)
    open_audio (p, p->oc, p->audio_st);

  if (url_fopen (&p->oc->pb, self->path, URL_WRONLY) < 0)
    {
      fprintf (stderr, "couldn't open '%s'\n", self->path);
      return -1;
    }

  av_write_header (p->oc);

  return 0;
}

#if 0
static int
filechanged (GeglChantOperation *op, const char *att)
{
  init (op);
  return 0;
}
#endif

static gboolean
process (GeglOperation       *operation,
         GeglBuffer          *input,
         const GeglRectangle *result,
         gint                 level)
{
  static gint inited = 0;
  GeglChantO *o = GEGL_CHANT_PROPERTIES (operation);
  Priv       *p = (Priv*)o->chant_data;

  g_assert (input);

  if (p == NULL)
    init (o);
  p = (Priv*)o->chant_data;

  p->width = result->width;
  p->height = result->height;
  p->input = input;

  if (!inited)
    {
      tfile (o);
      inited = 1;
    }

  write_video_frame (o, p->oc, p->video_st);
  if (p->audio_st)
    write_audio_frame (o, p->oc, p->audio_st);

  return  TRUE;
}

static void
finalize (GObject *object)
{
  GeglChantO *o = GEGL_CHANT_PROPERTIES (object);
  if (o->chant_data)
    {
      Priv *p = (Priv*)o->chant_data;

    if (p->oc)
      {
        gint i;
        if (p->video_st)
          close_video (p, p->oc, p->video_st);
        if (p->audio_st)
          close_audio (p, p->oc, p->audio_st);

        av_write_trailer (p->oc);

        for (i = 0; i < p->oc->nb_streams; i++)
          {
            av_freep (&p->oc->streams[i]);
          }

        url_fclose (&p->oc->pb);
        free (p->oc);
      }
      g_free (o->chant_data);
      o->chant_data = NULL;
    }

  G_OBJECT_CLASS (g_type_class_peek_parent (G_OBJECT_GET_CLASS (object)))->finalize (object);
}


static void
gegl_chant_class_init (GeglChantClass *klass)
{
  GeglOperationClass     *operation_class;
  GeglOperationSinkClass *sink_class;

  G_OBJECT_CLASS (klass)->finalize = finalize;

  operation_class = GEGL_OPERATION_CLASS (klass);
  sink_class      = GEGL_OPERATION_SINK_CLASS (klass);

  sink_class->process = process;
  sink_class->needs_full = TRUE;

  gegl_operation_class_set_keys (operation_class,
    "name"        , "gegl:ff-save",
    "categories"  , "output:video",
    "description" , _("FFmpeg video output sink"),
    NULL);
}

#endif
