//
//  VideoDecoder.m
//  DecoderWrapper
//
//  Copyright 2010 Dropcam. All rights reserved.
//

#import "VideoDecoder.h"

#import <avformat.h>
#import <avcodec.h>
#import <swscale.h>
#import "NLCommonLoggingNVP.h"

//#define SHOW_DEBUG_MV

#define OUTPUT_COLOR_FORMAT PIX_FMT_RGB555BE

LogCallbackfn g_logCallbackFn = NULL;

const uint8_t startCode[] = { 0x00, 0x00, 0x00, 0x01 };
#define INIT_INPUT_BUFFER_SIZE 2048

static uint8_t get_h264_nalu_type(NSData *packet)
{
  if ([packet length] < 5)
    return -1;
  
  uint8_t header[5] = {0};
  [packet getBytes:header length:sizeof(header)];
  
  // Assuming all packets begin with 4-byte marker
  uint8_t nal_unit_type = (header[4] & 0x1F);
  return nal_unit_type;
}

static void av_log_callback(void *ptr, 
                            int level, 
                            const char *fmt, 
                            va_list vl)
{
  static char line[1024] = {0};
  const char *module = "unknown";
  
  if (ptr)
  {
    AVClass *avc = *(AVClass**) ptr;
    module = avc->item_name(ptr);
  }
  
  vsnprintf(line, sizeof(line), fmt, vl);
  
  if (g_logCallbackFn) {
    g_logCallbackFn(level, module, line);
  }
}

@implementation VideoDecoder {
  struct AVCodec *codec;
  struct AVCodecContext *codecCtx;
  struct AVFrame *dstFrame;
  struct SwsContext *convertCtx;
  uint8_t *outputBuf;
  int outputBufLen;
  
  BOOL outputInit;
  BOOL frameReady;
  
  uint8_t *inputBuf;
  int inputBufLen;
  
  AVFrame *_srcFrame;
  AVFrame *_dstFrame;
  AVCodecContext *_codecCtx;
  enum AVPixelFormat _outputColorSpace;
}

+ (void)staticInitialize {
  av_register_all();  
}

+ (void)registerLogCallback:(LogCallbackfn)fn {
  g_logCallbackFn = fn;
  av_log_set_level(AV_LOG_DEBUG);  
  av_log_set_callback(av_log_callback);  
}

- (id)initWithCodec:(enum VideoCodecType)codecType 
         colorSpace:(enum VideoColorSpace)colorSpace 
              width:(int)width 
             height:(int)height 
        privateData:(NSData*)privateData
        allowChunks:(BOOL)allowChunks {
  if ((self = [super init])) {
    
    inputBufLen = INIT_INPUT_BUFFER_SIZE;
    inputBuf = malloc(inputBufLen);
    
    if (colorSpace == VideoColorSpaceRGB) {
      _outputColorSpace = PIX_FMT_RGB555BE;
    } else {
      _outputColorSpace = PIX_FMT_YUV420P;
    }
    
    codec = avcodec_find_decoder(CODEC_ID_H264);
    _codecCtx = avcodec_alloc_context3(codec);
    
    _codecCtx->thread_count = 0;
    _codecCtx->thread_type = FF_THREAD_FRAME;
    
    // Note: for H.264 RTSP streams, the width and height are usually not specified (width and height are 0).  
    // These fields will become filled in once the first frame is decoded and the SPS is processed.
    _codecCtx->width = width;
    _codecCtx->height = height;
    
    if (privateData != nil) {
      _codecCtx->extradata = av_malloc([privateData length]);
      _codecCtx->extradata_size = (int)[privateData length];
      [privateData getBytes:_codecCtx->extradata length:_codecCtx->extradata_size];
    }
    
    _codecCtx->pix_fmt = PIX_FMT_YUV420P;
#ifdef SHOW_DEBUG_MV
    _codecCtx->debug_mv = 0xFF;
#endif
    
    _codecCtx->flags |= CODEC_FLAG_EMU_EDGE;
    _codecCtx->flags2 |= CODEC_FLAG2_FAST;
    
    if (allowChunks) {
      _codecCtx->flags2 |= CODEC_FLAG2_CHUNKS;
    }
    
    _srcFrame = av_frame_alloc();
    _dstFrame = av_frame_alloc();

    // Lock required around calls to avcodec_open/close
    @synchronized(self.class) {
      int res = avcodec_open2(_codecCtx, codec, NULL);

      if (res < 0) {
        NLLogNVPWarn(@"Failed to initialize decoder");
      }
    }
  }
  
  return self;  
}

- (void)decodeFrame:(NSData*)frameData hasStartCode:(BOOL)hasStartCode PTS:(int64_t)PTS {
  AVPacket packet = {0};
  
  if (frameData) {
    int neededSize = (int)[frameData length] + 4 + (FF_INPUT_BUFFER_PADDING_SIZE);
    if (inputBufLen < neededSize) {
      inputBufLen = neededSize * 2;
      inputBuf = realloc(inputBuf, inputBufLen);
    }
    
    packet.data = inputBuf;

    if (hasStartCode) {
      memcpy(inputBuf, [frameData bytes], [frameData length]);
      packet.size = (int)[frameData length];
    } else {
      memcpy(inputBuf, startCode, 4);
      memcpy(inputBuf + 4, [frameData bytes], [frameData length]);
      packet.size = (int)[frameData length] + 4;
    }
  }
  
  packet.pts = PTS;
  
  int frameFinished = 0;
  int res = avcodec_decode_video2(_codecCtx, _srcFrame, &frameFinished, &packet);
  if (res < 0) {
#ifdef DEBUG
    NLLogNVPInfo(@"avcodec_decode_video2 res: %d", res);
#endif
    return;
  }
  
  if (frameFinished)
    frameReady = YES;
}

- (void)initializeOutput {
// Need to delay initializing the output buffers because we don't know the dimensions until we decode the first frame.
  if (!outputInit && _codecCtx->width > 0 && _codecCtx->height > 0) {
    outputBufLen = avpicture_get_size(_outputColorSpace, _codecCtx->width, _codecCtx->height);
    outputBuf = av_malloc(outputBufLen);
    
    avpicture_fill((AVPicture*)_dstFrame, outputBuf, _outputColorSpace, _codecCtx->width, _codecCtx->height);
    
    convertCtx = sws_getContext(_codecCtx->width, _codecCtx->height, _codecCtx->pix_fmt,  _codecCtx->width,
                                _codecCtx->height, _outputColorSpace, SWS_FAST_BILINEAR, NULL, NULL, NULL);
    
    outputInit = YES;
  }
}

- (BOOL)isFrameReady {
  return frameReady;
}

- (UIImage *)getDecodedRGBFrame {
  if (!frameReady)
    return nil;
  
  [self initializeOutput];
    
  sws_scale(convertCtx, (const uint8_t**)_srcFrame->data, _srcFrame->linesize, 0, _codecCtx->height, _dstFrame->data, _dstFrame->linesize);
  
  frameReady = NO;

  NSData *imgData = [NSData dataWithBytesNoCopy:outputBuf length:outputBufLen freeWhenDone:NO];
  
  NSUInteger width = [self getDecodedFrameWidth];
  NSUInteger height = [self getDecodedFrameHeight];
  
  if (width && height && imgData) {
    CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
    //The constants for specifying the alpha channel information are declared with the CGImageAlphaInfo type but can be passed to this parameter safely.
    CGContextRef context = CGBitmapContextCreate((void*)[imgData bytes], width, height, 5, width*2, colorSpace, (CGBitmapInfo)kCGImageAlphaNoneSkipFirst);
    
    CGImageRef imageRef = CGBitmapContextCreateImage(context);
    UIImage *result = [UIImage imageWithCGImage:imageRef];
    
    CGImageRelease(imageRef);
    CGContextRelease(context);
    CGColorSpaceRelease(colorSpace);
    
    return result;
  }
  
  return nil;
}

- (DecodedFrame *)getPixelBufferFromFrame:(struct AVFrame *)frame {
  NSUInteger width = [self getDecodedFrameWidth];
  NSUInteger height = [self getDecodedFrameHeight];
  
  NSData *yPlane = [NSData dataWithBytesNoCopy:frame->data[0] length:frame->linesize[0]*height freeWhenDone:NO];
  NSData *uPlane = [NSData dataWithBytesNoCopy:frame->data[1] length:frame->linesize[1]*height/2 freeWhenDone:NO];
  NSData *vPlane = [NSData dataWithBytesNoCopy:frame->data[2] length:frame->linesize[2]*height/2 freeWhenDone:NO];
  
  DecodedFrame *decodedFrame = [[DecodedFrame alloc] initWithWidth:width height:height yPlane:yPlane uPlane:uPlane vPlane:vPlane timestamp:0];
  
  return decodedFrame;
}

- (DecodedFrame *)getDecodedFrame {
  if (!frameReady)
    return nil;
  
  frameReady = NO;
  
  DecodedFrame *frame = nil;
  
  // Run it through sws_scale if the image is not packed
  if (_srcFrame->linesize[0] != _srcFrame->width) {
    [self initializeOutput];
    sws_scale(convertCtx, (const uint8_t**)_srcFrame->data, _srcFrame->linesize, 0, _codecCtx->height, _dstFrame->data, _dstFrame->linesize);
    frame = [self getPixelBufferFromFrame:_dstFrame];
  } else {
    // Otherwise we can just return the source frame
    frame = [self getPixelBufferFromFrame:_srcFrame];
  }

  if (_srcFrame->pkt_pts != AV_NOPTS_VALUE && _srcFrame->pkt_pts != 0) {
    frame.PTS = _srcFrame->pkt_pts;
  }
  
  return frame;
}

- (NSUInteger)getDecodedFrameWidth {
  return _codecCtx->width;
}

- (NSUInteger)getDecodedFrameHeight {
  return _codecCtx->height;
}

- (void)dealloc {
  free(inputBuf);
  
  // Lock required around calls to avcodec_open/close
  @synchronized(self.class) {
    av_free(_codecCtx->extradata);
    avcodec_close(_codecCtx);
    av_free(_codecCtx);
    av_frame_free(&_srcFrame);
    av_frame_free(&_dstFrame);
    av_free(outputBuf);
    sws_freeContext(convertCtx);
  }
}

@end