/*******************************************************************************
Copyright Datapath Ltd. 2014.

File:    rgbh264nal.cpp

History: 30 OCT 14    RL   Created.
         12 OCT 15    DC   Added further hard-coded URLs.
                           Included GOP Length as an option in URLs.
         02 MAR 16    DC   Round up frame rate to nearest Hz (100mHz).
                           Makes the exported RTSP URL friendlier to the
                           end user.

*******************************************************************************/

#define _CRT_SECURE_NO_WARNINGS
#include <Windows.h>
#include <WinError.h>
#include <stdio.h>
#include <tchar.h>

#include <api.h>

#include <dgctypes.h>
#include <dgcmedia.h>
#include <dgcmediahelper.h>

#include <rgb.h>
#include <rgbapi.h>
#include <rgberror.h>
#include <rgbh264nal.h>

/******************************************************************************/

uint32_t _Round(uint32_t value, uint32_t factor)
{
   value += (factor/2);
   value /= factor;
   value *= factor;
   return value;
}

/******************************************************************************/

uint32_t GetInputSignalType (
   uint32_t input,
   uint32_t *pWidth,
   uint32_t *pHeight,
   uint32_t *pFPS )
{
   unsigned long error;
   HRGBDLL  hDLL;

   /* Load the RGBEASY API. */
   error = RGBLoad(&hDLL);
   if (!error)
   {
      SIGNALTYPE signalType;

      error = RGBGetInputSignalType ( input, &signalType,
            (unsigned long*)pWidth, (unsigned long*)pHeight, (unsigned long*)pFPS );
      if (!error)
      {
         if (( signalType == RGB_SIGNALTYPE_NOSIGNAL) || 
            ( signalType == RGB_SIGNALTYPE_OUTOFRANGE ))
         {
            // use defaults
            error = RGBERROR_INVALID_FORMAT;
         }
         else
         {
            /* Round Framerate to nearest Hz (1000mHz) */
            *pFPS = _Round(*pFPS, 1000);
         }
      }
      RGBFree(hDLL);
   }
   return error;
}

/******************************************************************************/

uint32_t GetSupportedH264Inputs (
   uint32_t **ppInputList,
   uint32_t *pInputCount,
   uint32_t *pH264Count )
{
   unsigned long error;
   HRGBDLL  hDLL;

   /* Load the RGBEASY API. */
   error = RGBLoad(&hDLL);
   if (!error)
   {
      long bSupported = FALSE;
      uint32_t numberOfInputs;

      error = RGBGetNumberOfInputs ((unsigned long*)&numberOfInputs );
      if (!error)
      {
         uint32_t x, h264Count;
         uint32_t *pSupported = (uint32_t*)malloc (sizeof(uint32_t)*numberOfInputs);     

         h264Count=0;
         memset ( pSupported, 0, sizeof(uint32_t)*numberOfInputs);
         for ( x=0; x< numberOfInputs; x++)
         {
            error = RGBInputIsMediaSampleTypeSupported( x,
                                 DGCMEDIASAMPLETYPE_ENC_VIDEO,
                                 DGCENCVIDEOSAMPLESUBTYPE_H264,
                                 &bSupported );

            if ( !error && bSupported )
            {
               h264Count++;
               pSupported[x] = 1;
            }
         }
         if ( !h264Count )
            error = RGBERROR_ILLEGAL_CALL;
         else
         {
            *ppInputList = pSupported;
            *pInputCount = numberOfInputs;
            *pH264Count = h264Count;
         }
      }

      RGBFree(hDLL);
   }
   return error;
}

/******************************************************************************/

uint8_t GetNALType(uint8_t *pNalUnit) 
{
	return pNalUnit[0]&0x1F;
}

/******************************************************************************/

BOOL CanH264NALUnitBePurged(uint8_t *pNalUnit) 
{
	uint8_t nal_unit_type = pNalUnit[0]&0x1F;
	uint8_t nal_unit_data = pNalUnit[1];

	if ( nal_unit_type != 1 ) 
   { 
      return FALSE;
	}

   if( nal_unit_data == 0x88 )
   {
      return FALSE;
   }

	return TRUE;
}

/******************************************************************************/

BOOL DoesH264NALUnitBeginNewAccessUnit(uint8_t *pNalUnit) 
{
	uint8_t nal_unit_type = pNalUnit[0]&0x1F;

   // Non- VCL
	if ( ( nal_unit_type >= 6 && nal_unit_type <= 9 ) || 
        ( nal_unit_type >= 14 && nal_unit_type <= 18 )) 
           return TRUE;

	if ( nal_unit_type <= 5 && nal_unit_type > 0 ) 
   { 
      // We are a VCL NAL. Look at the high bit of the next byte;
      // if it's set, then this NAL starts a new 'access unit':
		if (( pNalUnit[1]&0x80) != 0 ) 
         return TRUE;
	}

	return FALSE;
}

/******************************************************************************/

CRGBEasyH264::CRGBEasyH264(uint32_t input) :
   m_error(RGBERROR_NO_ERROR),
   m_state(RGBEasyH264_FRAME_STATE_UNINITIALISED),
   m_hInstance(NULL),
   m_startCodeOffset(4),
   m_hDLL(NULL),
   m_input(input),
   m_hRGB(NULL),
   m_pNAL_List(NULL),
   m_hListMutex(NULL),
   m_queuedNALs(0),
   m_unqueuedNALs(0),
   m_NAL_ListSize(0)
{   
   m_hListMutex = CreateMutex ( NULL, FALSE, NULL );
   m_AdderMutex = CreateMutex ( NULL, FALSE, NULL );
   memset ( &m_StreamInfo, 0, sizeof(m_StreamInfo));
   m_error = RGBEasyH264Init( input );
}

CRGBEasyH264::~CRGBEasyH264()
{
   RGBEasyH264UnInit();
   CloseHandle(m_hListMutex);
   CloseHandle(m_AdderMutex);
}


/******************************************************************************/

PDGCMEDIASAMPLE CRGBEasyH264::CreateMediaSample(uint32_t width, uint32_t height)
{
   DGCMEDIASAMPLETYPE majorType = DGCMEDIASAMPLETYPE_ENC_VIDEO;
   DGCMEDIASAMPLESUBTYPE subType = DGCENCVIDEOSAMPLESUBTYPE_H264;

   PDGCMEDIASAMPLE pMediaSampleIn = MediaSampleAllocate( majorType, subType );
   if ( pMediaSampleIn )
   {
      if (MediaSampleAllocateFormatHeader(pMediaSampleIn))
      {
         PDGCVIDEOHEADER pVideoHeaderOut;

         // Scaled width and height of the capture
         pVideoHeaderOut = (PDGCVIDEOHEADER)pMediaSampleIn->PFormatHeader;
         pVideoHeaderOut->Width = width;
         pVideoHeaderOut->Height = height;

         if ( MediaSampleAllocateBufferHeader(
            pMediaSampleIn, DGCBUFFERHEADERTYPE_MEMORY ) )
         {
            if ( MediaSampleInitialisePlane(
               pMediaSampleIn, 0, width * height ) )
            {
               return pMediaSampleIn;
            }
         }
      }
      MediaSampleFree( pMediaSampleIn );
   }
   return NULL;
}

/******************************************************************************/

void CRGBEasyH264::AddNALToList ( PH264NAL pH264NAL )
{     
   /* Maintain a FIFO of NALU's. */
   if ( pH264NAL )
   {
      PH264NAL pH264NALPrev, pH264NALThis;
      PH264NAL pH264NALPrevToToRemove, pH264NALToRemove;

      /* Aquire Mutex. */
      WaitForSingleObject ( m_hListMutex, INFINITE );

      /* Insert at end of list. */
      pH264NALPrev = NULL;
      pH264NALThis = m_pNAL_List;
      pH264NALToRemove = NULL;
      pH264NALPrevToToRemove = NULL;

      while ( pH264NALThis )
      {
         if(pH264NALToRemove == NULL)
         {
            if(CanH264NALUnitBePurged(pH264NALThis->PStart))
            {
               // NAL could be removed from the list if the list is full
               pH264NALToRemove = pH264NALThis;
               pH264NALPrevToToRemove = pH264NALPrev;
            }
         }
         pH264NALPrev = pH264NALThis;
         pH264NALThis = pH264NALThis->PNext;
      }

      if ( pH264NALPrev )
      {
         pH264NALPrev->PNext = pH264NAL;
      }
      else
      {
         m_pNAL_List = pH264NAL;
      }

      m_queuedNALs++;
      m_NAL_ListSize++;
      pH264NAL->PNext = NULL;

      if( m_NAL_ListSize > 100)
      {
         if( pH264NALToRemove )
         {
            LockedRemoveNALFromList( pH264NALToRemove, pH264NALPrevToToRemove );
            m_NAL_ListSize--;
         }
      }

      /* Release Mutex. */
      ReleaseMutex ( m_hListMutex );
   }
}

/******************************************************************************/

void CRGBEasyH264::LockedRemoveNALFromList( PH264NAL pH264NALToRemove, PH264NAL pH264NALPrevToToRemove )
{
   PH264NAL pH264NAL = m_pNAL_List;
   while ( pH264NAL )
   {
      if(pH264NAL == pH264NALToRemove)
      {
         if(pH264NALPrevToToRemove)
         {
            // Link previous to next
            pH264NALPrevToToRemove->PNext = pH264NAL->PNext;
         }
         else
         {
            // Assign list head
            m_pNAL_List = pH264NAL->PNext;
         }

         pH264NAL->PNext = NULL; // Stop someone traversing the list through nefarious means.

         // Tidy up
         delete[] pH264NAL->PStart;
         delete[] pH264NAL;

         break;
      }
      pH264NAL = pH264NAL->PNext;
   }
}

/******************************************************************************/

void CRGBEasyH264::CheckListRemoveFirstNALFromList( PH264NAL* ppH264NAL )
{
   PH264NAL pH264NAL = NULL;

   /* Remove the entry from the FIFO list. */
   WaitForSingleObject ( m_hListMutex, INFINITE );

   if ( m_pNAL_List )
   {
      /* Get the first item in the list. */
      pH264NAL = m_pNAL_List;

      /* New pointer to start of list. */
      m_pNAL_List = m_pNAL_List->PNext;
      pH264NAL->PNext = NULL; // Stop someone traversing the list through nefarious means.
      m_unqueuedNALs++;
      m_NAL_ListSize--;
   }

   /* Release Mutex. */
   ReleaseMutex ( m_hListMutex );

   *ppH264NAL = pH264NAL;
}

/******************************************************************************/

void RGBCBKAPI CRGBEasyH264::RGBEasyH264EncoderErrorFn ( HWND           hWnd,
                                                         HRGB           hRGB,
                                                         unsigned long  error,
                                                         ULONG_PTR      pUserData,
                                                         unsigned long  *pReserved)
{
   CRGBEasyH264 *pH264NAL = (CRGBEasyH264*)pUserData;
   pH264NAL->RGBEasyH264RealEncoderErrorFn ( hRGB, error );
}

void CRGBEasyH264::RGBEasyH264RealEncoderErrorFn( HRGB            hRGB,
                                                  unsigned long   error)
{
   TCHAR buffer[255];

   wsprintf (buffer, TEXT("CRGBEasyH264::RGBEasyH264RealEncoderErrorFn 0x%x\n"), error );
   _tprintf( buffer );
}

/******************************************************************************/

void RGBCBKAPI CRGBEasyH264::RGBEasyH264EncoderFrameFn(  HWND              hWnd,
                                                         HRGB              hRGB,
                                                         PDGCMEDIASAMPLE   pMediaSample,
                                                         ULONG_PTR         pUserData )
{
   CRGBEasyH264 *pH264NAL = (CRGBEasyH264*)pUserData;
   if ( pH264NAL )
      pH264NAL->RGBEasyH264RealEncoderFrameFn ( hRGB, pMediaSample );
}

unsigned char * CRGBEasyH264::GetNextNAL(unsigned char * pBuffer, uint32_t *pBufferLen, unsigned char ** pNAL, unsigned long * NALlength)
{
   uint32_t nallen = *pBufferLen;
   bool bSeekingEnd = false;
   bool bSingleNAL = false;
   bool bIDRFrame = false;
   bool bSPS = false;
   bool bPPS = false;
   unsigned char * pData = pBuffer, *pNALStart;
   uint8_t NALType = 0;
   uint32_t BufLen = *pBufferLen;
   while (BufLen != 0)
   {
      unsigned long *pHeader;
         
      pHeader = reinterpret_cast<unsigned long *>(pData);
      if (*pHeader == 0x01000000)
      {
         NALType = GetNALType(pData + 4);
         if (bSeekingEnd)
         {

            // NAL Ends.
            *NALlength = pData - pNALStart;
            *pNAL = pNALStart;
            BufLen -= ((pData - pNALStart) + 4);
            *pBufferLen = BufLen;
            return pData;
         }
         // NAL begins
         pNALStart = pData+4; // Skip header.
         bSeekingEnd = true;
         pData += 3; // We know we can optimise these bytes out of the search.
         BufLen -= 3; // And commensurately shorten the buffer.
      }
      if(NALType != 7 && NALType != 8)
      {
         pData += BufLen;
         BufLen = 0;
      }
      else
      {
         pData++;
         BufLen -= 1;
      }
   }

   if (bSeekingEnd)
   {
      *NALlength = pData - pNALStart;
      *pNAL = pNALStart;
      *pBufferLen = 0;
      return pData;
   }
   // Non NAL data in the buffer...
   return NULL;
}

void CRGBEasyH264::RGBEasyH264RealEncoderFrameFn(  HRGB              hRGB,
                                                   PDGCMEDIASAMPLE   pMediaSample )
{   
   PDGCMEMORYBUFFERHEADER pHeader = (PDGCMEMORYBUFFERHEADER) pMediaSample->PBufferHeader;
   uint32_t size = pHeader->Planes[0].ActualLength;
   unsigned char *pStart = (unsigned char *)pHeader->Planes[0].PBuffer;
   // We need to protect against any other buffer dissector from starting to add its buffer's NALs
   // to the (singleton) list of NALs until we've completely finished adding our NALs from this buffer.
   // With any luck, this thread will run before any other later thread starts to, although timestamps
   // *should* sort any interleaving out.  Even so, mutexing against other threads running whilst we're
   // processing this buffer seems like a good idea...

   WaitForSingleObject(m_AdderMutex, INFINITE);

   /* An RGBEasy buffer containing elementary H264 data can contain multiple NAL units. 
    * However, a single buffer will always contain a single VCR NAL and possible multiple
    * non-VCR NAL units. For this reason we split the buffer into NAL units in this 
    * function and assign the same time stamp to all NAL units. */
   uint64_t TimeStamp = pHeader->StartTime;
   
   {
      unsigned char *pBuffer = pStart, *pNALData;
      unsigned long length;
      BOOL bAdded = FALSE;

      // GetNextNAL side-effects the size; it is adjusted down by the size of the data consumed by the
      // NAL which is found.
      // pBuffer is unmodified, but the position in the buffer after the NAL found is returned, and it's
      // beholden on the caller to pass that pointer in to the next call to continue the search.
      // size will be 0 when the search is over, and NULL will be returned.
      // Note that if there is trailing data which isn't a NAL, then length can be non-zero, but the
      // function will return NULL.
      while (size != 0 && ((pBuffer = GetNextNAL(pBuffer, &size, &pNALData, &length)) != NULL))
      {
         // Found a NAL.  Add it to the list.
         PH264NAL pNAL = new H264NAL;
         pNAL->StartTime = TimeStamp;
         pNAL->Size = length;
         pNAL->PStart = new uint8_t[length];
         memcpy(pNAL->PStart, pNALData, length);
         pNAL->PNext = NULL;
         AddNALToList ( pNAL );
         bAdded = TRUE;
      }
   }
   // And release the buffer NAL content processing mutex.
   ReleaseMutex(m_AdderMutex);

   // Kick off the next data Acquisition from the encoder (if we're still in a capturing state).
   if ( m_state == RGBEasyH264_FRAME_STATE_STARTED )
   {
      m_error = RGBChainMediaSample( hRGB, pMediaSample );
   }
}


/******************************************************************************/

uint32_t CRGBEasyH264::RGBEasyH264Start( uint32_t width,
                                         uint32_t height,
                                         uint32_t frameRate,
                                         DGCENCLEVEL level,
                                         DGCENCPROFILE profile,
                                         uint32_t bitrate,
                                         uint32_t keyframeinterval )
{
   if ( m_error )
   {
      printf("CRGBEasyH264::RGBEasyH264Start: oops m_error(0x%x)\n", m_error);
      return m_error;
   }
   if ( m_state != RGBEasyH264_FRAME_STATE_INITIALISED )
   {
      printf("CRGBEasyH264::RGBEasyH264Start: (m_state != RGBEasyH264_FRAME_STATE_INITIALISED)\n");
      return RGBERROR_INVALID_STATE;
   }

   m_StreamInfo.Width = width;
   m_StreamInfo.Height = height;
   m_StreamInfo.FPS = frameRate;

   if ( !m_StreamInfo.Width || !&m_StreamInfo.Height || !m_StreamInfo.FPS )
   {
      printf("CRGBEasyH264::RGBEasyH264Start: ( !m_StreamInfo.Width || !&m_StreamInfo.Height || !m_StreamInfo.FPS )\n");
      m_error = RGBERROR_INVALID_FORMAT;
      return m_error;
   }

   m_error = RGBOpenInput( m_input, &m_hRGB );
   if ( !m_error )
   {
      uint32_t i;

      m_error = RGBSetCaptureRate( m_hRGB, m_StreamInfo.FPS );
      if ( m_error )
      {
         printf("CRGBEasyH264::RGBEasyH264Start: RGBSetCaptureRate\n");
         return m_error;
      }

      m_error = RGBSetEncodeLevel( m_hRGB, level );
      if ( m_error )
      {
         printf("CRGBEasyH264::RGBEasyH264Start: RGBSetEncodeLevel\n");
         return m_error;
      }

      m_error = RGBSetEncodeProfile( m_hRGB, profile );
      if ( m_error )
      {
         printf("CRGBEasyH264::RGBEasyH264Start: RGBSetEncodeProfile\n");
         return m_error;
      }

      m_error = RGBSetEncodeBitRate( m_hRGB, bitrate );
      if ( m_error )
      {
         printf("CRGBEasyH264::RGBEasyH264Start: RGBSetEncodeBitRate\n");
         return m_error;
      }
      
      // set every 90th frame to be a key frame
      m_error = RGBSetEncodeKeyFrameInterval( m_hRGB, keyframeinterval );
      if ( m_error )
      {
         printf("CRGBEasyH264::RGBEasyH264Start: RGBSetEncodeKeyFrameInterval\n");
         return m_error;
      }

      m_error = RGBSetErrorFn( m_hRGB, RGBEasyH264EncoderErrorFn, (ULONG_PTR)this );
      if ( !m_error )
      {
         m_error = RGBSetMediaSampleCapturedFn( m_hRGB, RGBEasyH264EncoderFrameFn, (ULONG_PTR)this );
         if ( !m_error )
         {
            for ( i = 0; i < NUM_SAMPLES; i++ )
            {
               m_pMediaSample[i] = CreateMediaSample(m_StreamInfo.Width, m_StreamInfo.Height);
               {
                  PDGCMEMORYBUFFERHEADER pHeader = (PDGCMEMORYBUFFERHEADER) m_pMediaSample[i]->PBufferHeader;
               }
               m_error = RGBChainMediaSample( m_hRGB, m_pMediaSample[i] );
            }
            if ( !m_error )
            {
               m_error = RGBUseOutputBuffers(m_hRGB, TRUE);
               if (!m_error)
               {
                  m_error = RGBStartCapture(m_hRGB);
                  if (!m_error)
                     m_state = RGBEasyH264_FRAME_STATE_STARTED;
                  else
                  {
                     printf("CRGBEasyH264::RGBEasyH264Start: RGBStartCapture m_error(0x%x)\n", m_error);
                  }
                  return m_error;
               }
               RGBUseOutputBuffers(m_hRGB, FALSE);
            }
            
            for ( i = 0; i < NUM_SAMPLES; i++ )
            {
               MediaSampleFree( m_pMediaSample[i] );
            }
         }
      }
      RGBCloseInput(m_hRGB);
   }
   printf("CRGBEasyH264::RGBEasyH264Start: Hmm m_error(0x%x)\n", m_error);
   return m_error;
}

/******************************************************************************/

uint32_t CRGBEasyH264::RGBEasyH264GetNAL ( 
   uint8_t  *pNAL, 
   uint32_t maxSize, 
   uint32_t *pTrunkSize,
   uint32_t *pCopySize,
   uint64_t *pTimeStamp )
{
   if ( m_error )
      return m_error;
   if ( m_state != RGBEasyH264_FRAME_STATE_STARTED )
      return RGBERROR_INVALID_STATE;
 
   PH264NAL pH264NAL = NULL;

   CheckListRemoveFirstNALFromList( &pH264NAL );
   if ( pH264NAL )
   {
      if ( pH264NAL->Size > maxSize )
      {
         *pTrunkSize = pH264NAL->Size - maxSize;
			*pCopySize = maxSize;
      }
      else
      {
         *pTrunkSize = 0;
			*pCopySize = pH264NAL->Size;
      }
      *pTimeStamp = pH264NAL->StartTime;
      memcpy ( pNAL, pH264NAL->PStart, *pCopySize );
      delete[] pH264NAL->PStart;
      delete[] pH264NAL;
   }
   else
   {
      return RGBERROR_BUFFER_NOT_VALID;
   }
   return m_error;
}

/******************************************************************************/

uint32_t CRGBEasyH264::RGBEasyH264Stop()
{
   uint32_t tick=0;
   if ( m_error )
      return m_error;
   if ( m_state != RGBEasyH264_FRAME_STATE_STARTED )
      return RGBERROR_INVALID_STATE;
   
   m_state = RGBEasyH264_FRAME_STATE_STOPING;
   m_error = RGBUseOutputBuffers(m_hRGB, FALSE);
   if ( !m_error )
   {
      m_error = RGBStopCapture(m_hRGB);  
      if ( !m_error )
      {
         m_error = RGBCloseInput(m_hRGB);
         {
            for ( int i = 0; i < NUM_SAMPLES; i++ )
            {
               MediaSampleFree( m_pMediaSample[i] );
            }

            m_state = RGBEasyH264_FRAME_STATE_STOPPED;
         }
      }
   }
   return m_error;
}

/******************************************************************************/

uint32_t CRGBEasyH264::RGBEasyH264Init( uint32_t input )
{
   if ( m_error )
      return m_error;
   if ( m_state != RGBEasyH264_FRAME_STATE_UNINITIALISED )
      return RGBERROR_INVALID_STATE;

   /* Load the RGBEASY API. */
   m_error = RGBLoad(&m_hDLL);
   if (!m_error)
   {
      uint32_t bSupported = FALSE;

      m_error = RGBInputIsMediaSampleTypeSupported( input,
                              DGCMEDIASAMPLETYPE_ENC_VIDEO,
                              DGCENCVIDEOSAMPLESUBTYPE_H264,
                              (long*)&bSupported );

      if ( !m_error && bSupported )
      {
         m_state = RGBEasyH264_FRAME_STATE_INITIALISED;
         return m_error;
      }
      else
         m_error = RGBERROR_INVALID_INPUT;

      RGBFree(m_hDLL);
   }
   return m_error;
}

/******************************************************************************/

uint32_t CRGBEasyH264::RGBEasyH264UnInit( )
{
   if ( m_error )
      return m_error;
   if ( m_state != RGBEasyH264_FRAME_STATE_STOPPED )
      return RGBERROR_INVALID_STATE;

   if ( m_hDLL )
      RGBFree( m_hDLL );

   m_state = RGBEasyH264_FRAME_STATE_UNINITIALISED;
   return 0;
}

/******************************************************************************/