#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <windows.h>
#include <psapi.h>
#include <time.h>
#include <assert.h>
#include "mfxvideo.h"

#define SHOWHDRS 1
#define SHOWMEM  0

static void _initEncParam(mfxVideoParam *par);

int main(int argc, char **argv)
{
  int useHW = 1;
  int i,j;
  char type;
  mfxIMPL impl;
  mfxStatus sts;
  mfxVersion ver;
  mfxSession session;
  mfxSyncPoint sync;
  mfxBitstream bs;
  mfxVideoParam par;
  mfxFrameAllocRequest request;
  mfxFrameSurface1 *surfaces;
  FILE *fid = NULL;
  HANDLE pHandle;
  PROCESS_MEMORY_COUNTERS memInfoStart;
  PROCESS_MEMORY_COUNTERS memInfoNow;
  time_t tic,toc;
  int elapsed = 0;
  
  // Usage "msdkMem -s" use software only
  // Usage "msdkMem -h" use hardware only
  // Usage "msdkMem" uses hardware
  if (argc > 1) useHW = (argv[1][1] != 's');
  if (argc > 2)
  {
    fid = fopen(argv[2],"wb");
    if (fid == NULL) 
    {
      fprintf(stderr, "Could not open %s\n", argv[2]);
      return(-1);
    }
  }
  ver.Major = 1;
  ver.Minor = 1;
  
  // Create MSDK session
  impl = (useHW == 1) ? MFX_IMPL_HARDWARE_ANY:MFX_IMPL_SOFTWARE;
  sts = MFXInit(impl,&ver, &session);
  if (sts != MFX_ERR_NONE)
  {
    fprintf(stderr,"Could not init session:%d\n", sts);
    return(-1);
  }
  else
  {
    MFXQueryVersion(session, &ver);
    MFXQueryIMPL(session,&impl);
    fprintf(stderr,"Using %s version %d.%d\n",
          (impl == MFX_IMPL_SOFTWARE) ? "SW":"HW",
            ver.Major, ver.Minor);
  }
  
  // Create encoder
  _initEncParam(&par);
  sts = MFXVideoENCODE_Init(session, &par);
  if (sts == MFX_WRN_PARTIAL_ACCELERATION)
  {
    fprintf(stderr, "WARNING: Full acceleration is not being used\n");
    sts = MFX_ERR_NONE;
  }
  if (sts != MFX_ERR_NONE)
  {
    fprintf(stderr,"Could not create encoder:%d\n",sts);
    return(-1);
  }

  // Allocate dummy memory resources
  memset(&request, 0, sizeof(mfxFrameAllocRequest));
  sts = MFXVideoENCODE_QueryIOSurf(session,
                                   &par,
                                   &request);
  if (sts != MFX_ERR_NONE)
  {
    fprintf(stderr,"Error determining required buffer size:%d\n", sts);
    return(-1);
  }
  else
  {
    fprintf(stderr,"%d buffers suggested\n",request.NumFrameSuggested);
  }
  surfaces = (mfxFrameSurface1 *) calloc(request.NumFrameSuggested, sizeof(mfxFrameSurface1));
  for (i = 0;i < request.NumFrameSuggested;i++)
  {
	  unsigned char *data;
	  data = (unsigned char *) _aligned_malloc((1920*1088*3)/2, 128);
	  surfaces[i].Data.Y = data;
	  surfaces[i].Data.UV = data + 1920*1088;
	  surfaces[i].Data.Pitch = 1920;
	  surfaces[i].Info = par.mfx.FrameInfo;	  
  }
  bs.MaxLength = 134*2048;
  bs.Data   = malloc(bs.MaxLength);
  i = 0;
  j = 0;

  // Encode frames until end of time
  fprintf(stderr,"Starting frame encode\n");
#if SHOWMEM
  fprintf(stderr,"Elapsed (s) DeltaMemory\n");
#endif
  pHandle = GetCurrentProcess();
  tic = time(NULL);
  while (1)
  {
    int fidx;

    // Show memory delta periodically
    toc = time(NULL);
    if ((toc - tic) >= elapsed)
    {
      int dt;
      int dm;
      if (elapsed <= 10)
      {
        GetProcessMemoryInfo(pHandle, &memInfoStart, sizeof(memInfoStart));
      }
      else
      {
        GetProcessMemoryInfo(pHandle, &memInfoNow, sizeof(memInfoNow));
        dt = (int) (toc - tic - 20);
        dm = memInfoNow.WorkingSetSize - memInfoStart.WorkingSetSize;
#if SHOWMEM
        fprintf(stderr,"%11d %11d\n",dt,dm);
#endif
      }
      elapsed += 10;
    }

    // Find a free buffer
    for (fidx = 0;fidx < request.NumFrameSuggested;fidx++) 
      if (surfaces[fidx].Data.Locked == 0) break;
    if (fidx == request.NumFrameSuggested)
    {
      fprintf(stderr,"Ran out of buffers?\n");
      return(-1);
    }

    // Fill buffer with  dummy data
	  memset(surfaces[fidx].Data.Y,j*16, (1920*1088*3)/2);
    j++;

    // Start one frame encode
    bs.DataOffset = 0;
    bs.DataLength = 0;
    do
    {
      sts = MFXVideoENCODE_EncodeFrameAsync(session,
                                            NULL,
                                            &surfaces[fidx],
                                            &bs,
                                            &sync);
    } while (sts == MFX_WRN_DEVICE_BUSY);

    if (sts == MFX_ERR_MORE_DATA) continue;

    if (sts == MFX_ERR_NONE)
    {
      // Wait for encode to finish
      do
      {
        sts = MFXVideoCORE_SyncOperation(session,
                                         sync,
                                         1);
      } while (sts == MFX_WRN_IN_EXECUTION);
      if (sts != MFX_ERR_NONE)
      {
        fprintf(stderr,"Sync failed with %d\n", sts);
        return(-1);
      }

      // Show frame information
      switch (bs.FrameType & 0xf)
      {
      case MFX_FRAMETYPE_I:
        type = 'I';
        break;
      case MFX_FRAMETYPE_P:
        type = 'P';
        break;
      case MFX_FRAMETYPE_B:
        type = 'B';
        break;
      default:
        type = '?';
        break;
      }
#if SHOWHDRS
      fprintf(stderr,"Frame %03d of type %c is %6d bytes\n", i, type,bs.DataLength);
      // Parse the headers until slice data	  
	  {
       int bytes = bs.DataLength;
       unsigned char *ptr;
       int vbv_buffer_size_value;
       int vbv_buffer_size_ext;	   
       ptr = &bs.Data[bs.DataOffset];	   
       while (bytes > 4)
	   {
          if (ptr[0] == 0 && ptr[1] == 0 && ptr[2] == 1)
		  {
            ptr +=3;
            bytes -= 3;
            if (ptr[0] >= 1 && ptr[0] <= 0xaf)
			{
				fprintf(stderr,"Slice data starting\n");
                break;
			}
            switch (ptr[0])
			{
			case 0:				
              fprintf(stderr,"Picture Start Code\n");
			  break;
			case 0xb2:
              fprintf(stderr,"User Start Code\n");
			  break;
			case 0xb3:
				{				
			  bytes--;
			  ptr++;
              fprintf(stderr,"Sequence Header Code\n");
//              fprintf(stderr,"hsz = %d\n", (ptr[0] << 4) | ((ptr[1] >> 4) & 0xf));
//              fprintf(stderr,"vsz = %d\n", ((ptr[1] & 0xf) << 8) | ptr[2]);
              // ptr[3] -> frame/aspect codes
              // ptr[4]|ptr[5]|2b ptr[6] -> bit_rate_value
              // 3rd bit ptr[6] marker
   		      vbv_buffer_size_value = (ptr[6] & 0x1f) << 5 | ((ptr[7] >> 3) & 0x1f);
				}
			  break;
			case 0xb4:
              fprintf(stderr,"Sequence Error Code\n");
			  break;
			case 0xb5:
              fprintf(stderr,"Extension Start Code\n");
			  bytes--;
			  ptr++;
              if (((ptr[0] >> 4) & 0xf) == 0x1)
			  {
                  unsigned char profile,level;				  
				  fprintf(stderr,"Sequence Extension\n");
                  profile = ptr[0] & 0x7;
                  level   = (ptr[1] >> 4) & 0xf;
                  switch (profile)
				  {
				  case 0x5:
                  fprintf(stderr,"Simple Profile\n");
				  break;
				  case 0x4:
                  fprintf(stderr,"Main Profile\n");
				  break;
				  case 0x3:
                  fprintf(stderr,"SNR Scalable\n");
				  break;
				  case 0x2:
                  fprintf(stderr,"Spatially Scalable\n");
				  break;
				  case 0x1:
                  fprintf(stderr,"High Profile\n");
				  break;
				  default:
                  fprintf(stderr,"Invalid Profile: %0x\n", profile);
				  break;
				  }
                  switch (level)
				  {
				  case 0xa:
                  fprintf(stderr,"Low Level\n");
				  break;
				  case 0x8:
                  fprintf(stderr,"Main Level\n");
				  break;
				  case 0x6:
                  fprintf(stderr,"High 1440 Level\n");
				  break;
				  case 0x4:
                  fprintf(stderr,"High Level\n");
				  break;
				  }
   		      vbv_buffer_size_ext = ptr[4];
              fprintf(stderr,"vbvsz = %d\n", vbv_buffer_size_value | (vbv_buffer_size_ext << 10));			  
              assert((vbv_buffer_size_value | (vbv_buffer_size_ext << 10)) <= 597);
			  }				  
			  break;
			case 0xb7:
              fprintf(stderr,"End Sequence Code\n");
			  break;
			case 0xb8:
              fprintf(stderr,"Group Start Code\n");
			  break;
			default:
              fprintf(stderr,"Unknown Start Code: %0x\n", ptr[0]);
			  break;
			}
		  }			
		  bytes--;
		  ptr++;
	   }	   
		  
	  }		  
#endif
      if (fid)
      {
        fwrite(&bs.Data[bs.DataOffset], bs.DataLength, 1, fid);
        fflush(fid);
      }
      i++;
    }
    else
    {
      fprintf(stderr,"MFXVideoENCODE_EncodeFrameAsync failed with %d\n", 
              sts);
    }
  }
  
  return(0);
}

// Configure encoder for MPEG2 with 1080i input, 16 Mb/s out
static void _initEncParam(mfxVideoParam *par)
{
  mfxExtCodingOption *option;
  mfxExtCodingOption **option_array;

  memset(par, 0, sizeof(mfxVideoParam));
  par->AsyncDepth = 0;
  par->IOPattern  = MFX_IOPATTERN_IN_SYSTEM_MEMORY;
  option =  (mfxExtCodingOption *) malloc(sizeof(mfxExtCodingOption));
  option_array = (mfxExtCodingOption**) malloc(sizeof(mfxExtCodingOption*));
  memset(option, 0, sizeof(mfxExtCodingOption));
  option->Header.BufferId                    = MFX_EXTBUFF_CODING_OPTION;
  option->Header.BufferSz                    = sizeof(mfxExtCodingOption);
  option->FramePicture                       = MFX_CODINGOPTION_ON;
  option_array[0] = option;
  par->ExtParam                    = (mfxExtBuffer **) option_array;
  par->NumExtParam                 = 1;
  par->mfx.BRCParamMultiplier      = 0;
  par->mfx.CodecId                 = MFX_CODEC_MPEG2;
  par->mfx.CodecProfile            = MFX_PROFILE_MPEG2_MAIN;
  par->mfx.CodecLevel              = MFX_LEVEL_MPEG2_MAIN;
  par->mfx.NumThread               = 0;
  par->mfx.TargetUsage             = MFX_TARGETUSAGE_BEST_QUALITY;
  par->mfx.TargetKbps              = 3000;
  par->mfx.GopPicSize              = 15;
  par->mfx.GopRefDist              = 3;
  par->mfx.GopOptFlag              = 0;
  par->mfx.IdrInterval             = 1;
  par->mfx.RateControlMethod       = (mfxU16) MFX_RATECONTROL_CBR;
  par->mfx.EncodedOrder            = 0;
  par->mfx.FrameInfo.FourCC       = MFX_FOURCC_NV12;
  par->mfx.FrameInfo.Width        = 720;
  par->mfx.FrameInfo.Height       = 480;
  par->mfx.FrameInfo.CropX = 0;
  par->mfx.FrameInfo.CropY = 0;
  par->mfx.FrameInfo.CropW = 720;
  par->mfx.FrameInfo.CropH = 480;
  par->mfx.FrameInfo.FrameRateExtD = 1001;
  par->mfx.FrameInfo.FrameRateExtN = 30000;
  par->mfx.FrameInfo.PicStruct    = MFX_PICSTRUCT_FIELD_TFF;
  par->mfx.FrameInfo.ChromaFormat = MFX_CHROMAFORMAT_YUV420;
}