/* This file is in the public domain. */

/*
 * This file exists because the codec include files are severely
 *  abusive of their caller's namespace; as one simple example, they
 *  #define fprintf please_use_av_log_instead_of_fprintf
 *  apparently for internal error-checking - but this gets imposed on
 *  clients of the interface too(!), because avcodec_open2, an
 *  interface intended for users to use, takes an AVDictionary pointer,
 *  and AVDictionary is defined in the file that does that #define.
 *
 * So we sequester all this grossness off in its own file.
 *  Unfortunately it is difficult to print anything here without
 *  undeffing avcodec.h's overrides - including even so much as
 *  implementing our av_log_set_callback argument!  I really have to
 *  wonder how this interface was designed to be used.
 */

#include <stdio.h>
#include <stdlib.h>
#include <strings.h>

#include "avcodec.h"
#include "avutil-mem.h"
#include "avutil-pixdesc.h"
#include "avutil-imgutils.h"

#include "codecintf.h"

extern AVCodec ff_h264_decoder;

typedef enum {
	  MT_AUDIO = 1,
	  MT_VIDEO,
	  } MEDIATYPE;

struct codec {
  MEDIATYPE type;
  unsigned char *extrabuf;
  union {
    struct {
      AVCodec *c;
      AVCodecContext *cc;
      AVFrame *frame;
      void (*framefn)(void *, int, int, const unsigned char *);
      void *ffnarg;
      AVPacket pkt;
      } v;
    } ;
  } ;

static unsigned char (*rgbdata)[3] = 0;
static int rgbwidth;
static int rgbheight;
static unsigned char *pktbuf = 0;
static int pktbuf_alloc = 0;
static int log_set_up = 0;
static int loglev;
static char *logbuf;
static int logalloc;
static int loglen;

#undef malloc
#undef free
#undef realloc
#undef exit
#undef printf
#undef fprintf

static void gen_rgb(CODEC *c)
{
 int y;
 int x;
 int ss;
 unsigned char *sp;
 unsigned char (*dp)[3];

 if (rgbdata) free(rgbdata);
 rgbwidth = c->v.cc->width;
 rgbheight = c->v.cc->height;
 rgbdata = malloc(rgbwidth*rgbheight*3);
 switch (c->v.cc->pix_fmt)
  { case PIX_FMT_YUV420P:
       ss = c->v.frame->linesize[0];
       sp = c->v.frame->data[0];
       dp = rgbdata;
       for (y=c->v.cc->height;y>0;y--)
	{ for (x=c->v.cc->width;x>0;x--)
	   { (*dp)[0] = *sp;
	     dp ++;
	     sp ++;
	   }
	  sp += ss - c->v.cc->width;
	}
       ss = c->v.frame->linesize[1];
       sp = c->v.frame->data[1];
       dp = rgbdata;
       for (y=c->v.cc->height>>1;y>0;y--)
	{ for (x=c->v.cc->width>>1;x>0;x--)
	   { (*dp)[1] = *sp;
	     dp ++;
	     (*dp)[1] = *sp;
	     dp ++;
	     sp ++;
	   }
	  sp -= c->v.cc->width >> 1;
	  for (x=c->v.cc->width>>1;x>0;x--)
	   { (*dp)[1] = *sp;
	     dp ++;
	     (*dp)[1] = *sp;
	     dp ++;
	     sp ++;
	   }
	  sp += ss - (c->v.cc->width >> 1);
	}
       ss = c->v.frame->linesize[2];
       sp = c->v.frame->data[2];
       dp = rgbdata;
       for (y=c->v.cc->height>>1;y>0;y--)
	{ for (x=c->v.cc->width>>1;x>0;x--)
	   { (*dp)[2] = *sp;
	     dp ++;
	     (*dp)[2] = *sp;
	     dp ++;
	     sp ++;
	   }
	  sp -= c->v.cc->width >> 1;
	  for (x=c->v.cc->width>>1;x>0;x--)
	   { (*dp)[2] = *sp;
	     dp ++;
	     (*dp)[2] = *sp;
	     dp ++;
	     sp ++;
	   }
	  sp += ss - (c->v.cc->width >> 1);
	}
       dp = rgbdata;
       for (y=c->v.cc->height;y>0;y--)
	{ for (x=c->v.cc->width;x>0;x--)
	   { int Y;
	     int U;
	     int V;
	     int R;
	     int G;
	     int B;
	     Y = (*dp)[0];
	     U = (*dp)[1] - 128;
	     V = (*dp)[2] - 128;
	     Y = (Y * 255) / 219;
	     U = (U * 255) / 224;
	     V = (V * 255) / 224;
	     R = (65536 * Y) + (91880 * V);
	     G = (65536 * Y) - (22580 * U) - (46799 * V);
	     B = (65536 * Y) + (116128 * U);
	     if (R > 0xffffff) R = 0xff; else R >>= 16;
	     if (G > 0xffffff) G = 0xff; else if (G < 0) G = 0; else G >>= 16;
	     if (B > 0xffffff) B = 0xff; else B >>= 16;
	     (*dp)[0] = R;
	     (*dp)[1] = G;
	     (*dp)[2] = B;
	     dp ++;
	   }
	}
       break;
    default:
       printf("*** don't know how to convert colour space %d (%s) to RGB\n",
		(int)c->v.cc->pix_fmt, av_pix_fmt_descriptors[(int)c->v.cc->pix_fmt].name);
       exit(1);
       break;
  }
}

static void logcb(void *arg __attribute__((__unused__)), int lev, const char *fmt, va_list ap)
{
 char *s;
 char *s0;
 int sl;
 char *nl;
 int n;
 int anylines;

 if ((loglen > 0) && (lev != loglev))
  { printf("*** log level changed %d -> %d\n",loglev,lev);
  }
 loglev = lev;
 anylines = 0;
 sl = vasprintf(&s,fmt,ap);
 s0 = s;
 while (sl > 0)
  { nl = memchr(s,'\n',sl);
    if (! nl)
     { if (loglen+sl > logalloc) logbuf = realloc(logbuf,logalloc=loglen+sl);
       bcopy(s,logbuf+loglen,sl);
       loglen += sl;
       break;
     }
    n = nl - s;
    printf("*** [%d] ",loglev);
    if (loglen > 0) fwrite(logbuf,1,loglen,stdout);
    fwrite(s,1,n+1,stdout);
    anylines = 1;
    sl -= n + 1;
    s += n + 1;
    loglen = 0;
  }
 if (anylines) fflush(stdout);
 free(s0);
}

static void video_decode_loop(CODEC *c)
{
 int l;
 int gotframe;

 bzero(pktbuf+c->v.pkt.size,FF_INPUT_BUFFER_PADDING_SIZE);
 while (c->v.pkt.size > 0)
  { l = avcodec_decode_video2(c->v.cc,c->v.frame,&gotframe,&c->v.pkt);
    if (l < 0)
     { printf("*** error decoding video\n");
       exit(1);
     }
    if (gotframe)
     { gen_rgb(c);
       (*c->v.framefn)(c->v.ffnarg,rgbwidth,rgbheight,&rgbdata[0][0]);
     }
    else if (l == 0)
     { printf("*** decoder returned 0 but no frame\n");
       exit(1);
     }
    c->v.pkt.size -= l;
    c->v.pkt.data += l;
  }
}

static void maybe_init_av_log(void)
{
 if (log_set_up) return;
 av_log_set_callback(&logcb);
 logbuf = 0;
 logalloc = 0;
 loglen = 0;
 log_set_up = 1;
}

static void pktbuf_ensure_size(int nb)
{
 if (nb > pktbuf_alloc)
  { free(pktbuf);
    pktbuf = malloc(pktbuf_alloc=nb);
  }
}

CODEC *codec_init_video(const unsigned char type[4], void (*cb)(void *, int, int, const unsigned char *), void *cbarg)
{
 CODEC *c;

 maybe_init_av_log();
 c = malloc(sizeof(CODEC));
 c->type = MT_VIDEO;
 c->extrabuf = 0;
 if (! bcmp(&type[0],"avc1",4))
  { c->v.c = &ff_h264_decoder;
  }
 else
  { printf("*** unknown video codec type\n");
  }
 c->v.cc = avcodec_alloc_context3(c->v.c);
 c->v.frame = avcodec_alloc_frame();
 c->v.framefn = cb;
 c->v.ffnarg = cbarg;
 av_init_packet(&c->v.pkt);
 return(c);
}

void codec_prefix(CODEC *c, const void *data, int len)
{
 c->extrabuf = malloc(len+FF_INPUT_BUFFER_PADDING_SIZE);
 bcopy(data,c->extrabuf,len);
 bzero(&c->extrabuf+len,FF_INPUT_BUFFER_PADDING_SIZE);
}

void codec_data(CODEC *c, const void *data, int len)
{
// cc->debug = ~0U;
 switch (c->type)
  { case MT_AUDIO:
       break;
    case MT_VIDEO:
       if (! avcodec_is_open(c->v.cc))
	{ if (c->extrabuf)
	   { c->v.cc->extradata = c->extrabuf;
	     c->v.cc->extradata_size = len;
	   }
	  if (avcodec_open2(c->v.cc,c->v.c,0) < 0)
	   { printf("*** codec open failed\n");
	     exit(1);
	   }
	}
       pktbuf_ensure_size(len+FF_INPUT_BUFFER_PADDING_SIZE);
       bcopy(data,pktbuf,len);
       c->v.pkt.size = len;
       c->v.pkt.data = pktbuf;
       video_decode_loop(c);
       break;
  }
}

void codec_done(CODEC *c)
{
 switch (c->type)
  { case MT_AUDIO:
       break;
    case MT_VIDEO:
       c->v.pkt.size = 0;
       c->v.pkt.data = 0;
       video_decode_loop(c);
       avcodec_close(c->v.cc);
       av_free(c->v.cc);
       av_free(c->v.frame);
       break;
  }
 free(c->extrabuf);
 free(c);
}