//////////////////////////////////////////////////////////////////////////
//
// preview.cpp: Manages video preview.
//
// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
// PARTICULAR PURPOSE.
//
// Copyright (c) Microsoft Corporation. All rights reserved.
//
//////////////////////////////////////////////////////////////////////////

#include "ARMFCaptureD3D.h"
#include "debug.h"
#include <shlwapi.h>
#include "ARAnalyse.h"

extern ARAnalyse	*g_pAnalyse;

extern HWND g_hStatus;

//-------------------------------------------------------------------
//  CreateInstance
//
//  Static class method to create the CPreview object.
//-------------------------------------------------------------------

HRESULT CPreview::CreateInstance(
    HWND hVideo,        // Handle to the video window.
    HWND hEvent,        // Handle to the window to receive notifications.
    CPreview **ppPlayer // Receives a pointer to the CPreview object.
    )
{
    assert(hVideo != NULL);
    assert(hEvent != NULL);

    if (ppPlayer == NULL)
    {
        return E_POINTER;
    }

    CPreview *pPlayer = new (std::nothrow) CPreview(hVideo, hEvent);

    // The CPlayer constructor sets the ref count to 1.

    if (pPlayer == NULL)
    {
        return E_OUTOFMEMORY;
    }

    HRESULT hr = pPlayer->Initialize();

    if (SUCCEEDED(hr))
    {
        *ppPlayer = pPlayer;
        (*ppPlayer)->AddRef();
    }

    SafeRelease(&pPlayer);
    return hr;
}


//-------------------------------------------------------------------
//  constructor
//-------------------------------------------------------------------

CPreview::CPreview(HWND hVideo, HWND hEvent) :
    m_pReader(NULL),
	m_pSource(NULL),    
    m_hwndVideo(hVideo),
    m_hwndEvent(hEvent),
    m_nRefCount(1),
    m_pwszSymbolicLink(NULL),
    m_cchSymbolicLink(0)
{
    InitializeCriticalSection(&m_critsec);
	m_video = true;
	m_transform = false;
	m_moyen = false;
}

//-------------------------------------------------------------------
//  destructor
//-------------------------------------------------------------------

CPreview::~CPreview()
{
    CloseDevice();
    m_draw.DestroyDevice();
	DeleteCriticalSection(&m_critsec);
}


//-------------------------------------------------------------------
//  Initialize
//
//  Initializes the object.
//-------------------------------------------------------------------

HRESULT CPreview::Initialize()
{
    HRESULT hr = S_OK;

    hr = m_draw.CreateDevice(m_hwndVideo);

    return hr;
}

bool CPreview::ToggleVideo()
{
	m_video = !m_video;
	return m_video;
}

bool CPreview::ToggleTransform()
{
	m_transform = !m_transform;
	return m_transform;
}
bool CPreview::ToggleMoyen()
{
	m_moyen = !m_moyen;
	return m_moyen;
}

//-------------------------------------------------------------------
//  CloseDevice
//
//  Releases all resources held by this object.
//-------------------------------------------------------------------

HRESULT CPreview::CloseDevice()
{
    EnterCriticalSection(&m_critsec);

	if ( m_pSource )
	{
		m_pSource->Stop();
		m_pSource->Shutdown();
	}

    SafeRelease(&m_pReader);
	SafeRelease(&m_pSource);
    

    CoTaskMemFree(m_pwszSymbolicLink);
    m_pwszSymbolicLink = NULL;
    m_cchSymbolicLink = 0;

    LeaveCriticalSection(&m_critsec);
    return S_OK;
}


/////////////// IUnknown methods ///////////////

//-------------------------------------------------------------------
//  AddRef
//-------------------------------------------------------------------

ULONG CPreview::AddRef()
{
    return InterlockedIncrement(&m_nRefCount);
}


//-------------------------------------------------------------------
//  Release
//-------------------------------------------------------------------

ULONG CPreview::Release()
{
    ULONG uCount = InterlockedDecrement(&m_nRefCount);
    if (uCount == 0)
    {
        delete this;
    }
    // For thread safety, return a temporary variable.
    return uCount;
}



//-------------------------------------------------------------------
//  QueryInterface
//-------------------------------------------------------------------

HRESULT CPreview::QueryInterface(REFIID riid, void** ppv)
{
    static const QITAB qit[] =
    {
        QITABENT(CPreview, IMFSourceReaderCallback),
        { 0 },
    };
    return QISearch(this, qit, riid, ppv);
}


/////////////// IMFSourceReaderCallback methods ///////////////

//-------------------------------------------------------------------
// OnReadSample
//
// Called when the IMFMediaSource::ReadSample method completes.
//-------------------------------------------------------------------

HRESULT CPreview::OnReadSample(
    HRESULT hrStatus,
    DWORD dwStreamIndex,
    DWORD /*dwStreamFlags*/,
    LONGLONG llTimestamp,
    IMFSample *pSample      // Can be NULL
    )
{
    HRESULT hr = hrStatus;
    IMFMediaBuffer *pBuffer = NULL;
	DWORD cbCurrentLength = 0;
	
    EnterCriticalSection(&m_critsec);
 
    if (SUCCEEDED(hr))
    {
        if (pSample)
        {
		// Get the video frame buffer from the sample.
		hr = pSample->ConvertToContiguousBuffer(&pBuffer);

		// AR analyse
		if ( g_pAnalyse && SUCCEEDED(hr) )
		{
			BYTE *image_buffer=0;
			pBuffer->Lock( &image_buffer, 0, &cbCurrentLength );
			double time = (double)(llTimestamp)/ 10000000.0;
			// Gray scale & blur & threshold
			if ( m_transform) 
				g_pAnalyse->Transform( image_buffer, cbCurrentLength, m_moyen );
			g_pAnalyse->Analyse(time, image_buffer, cbCurrentLength );
			pBuffer->Unlock();
		}
		// Draw the frame.

		if ( m_video && SUCCEEDED(hr))
		{
			hr = m_draw.DrawFrame(pBuffer);
		}

		}
    }
	// Request the next frame.
    if (m_pReader && SUCCEEDED(hr))
    {
        hr = m_pReader->ReadSample(
            dwStreamIndex,
            0,
            NULL,   // actual
            NULL,   // flags
            NULL,   // timestamp
            NULL    // sample
            );
    }

    if (FAILED(hr))
    {
        NotifyError(hr);
    }
    if ( pBuffer ) SafeRelease(&pBuffer);

    LeaveCriticalSection(&m_critsec);

	if (g_pAnalyse && SUCCEEDED(hr))
	{
		// display status message 
		WCHAR szBuf[1024];
		
		StringCbPrintf(szBuf, sizeof(szBuf),TEXT("Markers count=%3d confidence=( %0.1f / %0.1f / %0.1f ) FPS=( %0.1f / %0.1f / %0.1f ) delay=( %0.3f / %0.3f / %0.3f )"),
			g_pAnalyse->markers.size(), 
			g_pAnalyse->confidence.Min(), g_pAnalyse->confidence.Max(), g_pAnalyse->confidence.Avg() ,
			g_pAnalyse->frame_per_seconds.Min(), g_pAnalyse->frame_per_seconds.Max(), g_pAnalyse->frame_per_seconds.Avg(),
			g_pAnalyse->frame_delay.Min(), g_pAnalyse->frame_delay.Max(), g_pAnalyse->frame_delay.Avg() );
		SendMessage(g_hStatus, SB_SETTEXT, 0, (LPARAM)(LPSTR)szBuf );
	}
    
    return hr;
}


//-------------------------------------------------------------------
// TryMediaType
//
// Test a proposed video format.
//-------------------------------------------------------------------

HRESULT CPreview::TryMediaType(IMFMediaType *pType)
{
    HRESULT hr = S_OK;

    BOOL bFound = FALSE;
    GUID subtype = { 0 };
	
    hr = pType->GetGUID(MF_MT_SUBTYPE, &subtype);

    if (FAILED(hr))
    {
        return hr;
    }

    // Do we support this type directly?
    if (m_draw.IsFormatSupported(subtype))
    {
        bFound = TRUE;
    }
    else
    {
        // Can we decode this media type to one of our supported
        // output formats?

        for (DWORD i = 0;  ; i++)
        {
            // Get the i'th format.
            hr = m_draw.GetFormat(i, &subtype);
			if (FAILED(hr)) { break; }

			hr = pType->SetGUID(MF_MT_SUBTYPE, subtype);

            if (FAILED(hr)) { break; }

            // Try to set this type on the source reader.
            hr = m_pReader->SetCurrentMediaType(
                (DWORD)MF_SOURCE_READER_FIRST_VIDEO_STREAM,
                NULL,
                pType
                );

            if (SUCCEEDED(hr))
            {
                bFound = TRUE;
                break;
            }
        }
    }

    if (bFound)
    {
		EnterCriticalSection(&m_critsec);
        hr = m_draw.SetVideoType(pType);
		LeaveCriticalSection(&m_critsec);
    }

    return hr;
}

//-------------------------------------------------------------------
// SetDevice
//
// Set up preview for a specified video capture device.
//-------------------------------------------------------------------

HRESULT CPreview::SetDevice(IMFActivate *pActivate, int formatType)
{
    HRESULT hr = S_OK;

    IMFAttributes   *pAttributes = NULL;
    IMFMediaType    *pType = NULL;

    EnterCriticalSection(&m_critsec);

    // Release the current device, if any.

    hr = CloseDevice();

    // Create the media source for the device.
    if (SUCCEEDED(hr))
    {
        hr = pActivate->ActivateObject(
            __uuidof(IMFMediaSource),
            (void**)&m_pSource
            );
    }

    // Get the symbolic link.
    if (SUCCEEDED(hr))
    {
        hr = pActivate->GetAllocatedString(
            MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_VIDCAP_SYMBOLIC_LINK,
            &m_pwszSymbolicLink,
            &m_cchSymbolicLink
            );
    }
	
    //
    // Create the source reader.
    //
	
    // Create an attribute store to hold initialization settings.

    if (SUCCEEDED(hr))
    {
        hr = MFCreateAttributes(&pAttributes, 2);
    }
    if (SUCCEEDED(hr))
    {
        hr = pAttributes->SetUINT32(MF_READWRITE_DISABLE_CONVERTERS, TRUE);
    }

    // Set the callback pointer.
    if (SUCCEEDED(hr))
    {
        hr = pAttributes->SetUnknown(
            MF_SOURCE_READER_ASYNC_CALLBACK,
            this
            );
    }

    if (SUCCEEDED(hr))
    {
        hr = MFCreateSourceReaderFromMediaSource(
            m_pSource,
            pAttributes,
            &m_pReader
            );
    }


    // Try to find a suitable output type.
    if (SUCCEEDED(hr))
    {
		
        for (DWORD i = formatType; ; i++)
        {
            hr = m_pReader->GetNativeMediaType(
                (DWORD)MF_SOURCE_READER_FIRST_VIDEO_STREAM,
                i,
                &pType
                );

            if (FAILED(hr)) { break; }

            //hr = TryMediaType(pType); already done in SetFormat !!
			hr = SetFormat(pType);
			
			SafeRelease(&pType);

            if (SUCCEEDED(hr))
            {
                // Found an output type.
				
                break;
            }
        }
    }
	if (SUCCEEDED(hr))
    {
        // Ask for the first sample.
        hr = m_pReader->ReadSample(
            (DWORD)MF_SOURCE_READER_FIRST_VIDEO_STREAM,
            0,
            NULL,
            NULL,
            NULL,
            NULL
            );
    }

    if (FAILED(hr))
    {
        if (m_pSource)
        {
            m_pSource->Shutdown();

            // NOTE: The source reader shuts down the media source
            // by default, but we might not have gotten that far.
        }
        CloseDevice();
    }


    SafeRelease(&pAttributes);
    SafeRelease(&pType);

    LeaveCriticalSection(&m_critsec);
    return hr;
}



//-------------------------------------------------------------------
//  ResizeVideo
//  Resizes the video rectangle.
//
//  The application should call this method if the size of the video
//  window changes; e.g., when the application receives WM_SIZE.
//-------------------------------------------------------------------

HRESULT CPreview::ResizeVideo(WORD /*width*/, WORD /*height*/)
{
    HRESULT hr = S_OK;

    EnterCriticalSection(&m_critsec);

    hr = m_draw.ResetDevice();

    if (FAILED(hr))
    {
        MessageBox(NULL, L"ResetDevice failed!", NULL, MB_OK);
    }

    LeaveCriticalSection(&m_critsec);

    return hr;
}


//-------------------------------------------------------------------
//  CheckDeviceLost
//  Checks whether the current device has been lost.
//
//  The application should call this method in response to a
//  WM_DEVICECHANGE message. (The application must register for
//  device notification to receive this message.)
//-------------------------------------------------------------------

HRESULT CPreview::CheckDeviceLost(DEV_BROADCAST_HDR *pHdr, BOOL *pbDeviceLost)
{
    DEV_BROADCAST_DEVICEINTERFACE *pDi = NULL;

    if (pbDeviceLost == NULL)
    {
        return E_POINTER;
    }

    *pbDeviceLost = FALSE;

    if (pHdr == NULL)
    {
        return S_OK;
    }

    if (pHdr->dbch_devicetype != DBT_DEVTYP_DEVICEINTERFACE)
    {
        return S_OK;
    }

    pDi = (DEV_BROADCAST_DEVICEINTERFACE*)pHdr;


    EnterCriticalSection(&m_critsec);

    if (m_pwszSymbolicLink)
    {
        if (_wcsicmp(m_pwszSymbolicLink, pDi->dbcc_name) == 0)
        {
            *pbDeviceLost = TRUE;
        }
    }

    LeaveCriticalSection(&m_critsec);

    return S_OK;
}

HRESULT CPreview::EnumerateCaptureFormats(IMFMediaType ***pppTypes, UINT32 *pcCount)
{
    IMFPresentationDescriptor *pPD = NULL;
    IMFStreamDescriptor *pSD = NULL;
    IMFMediaTypeHandler *pHandler = NULL;
    IMFMediaType *pType = NULL;

    HRESULT hr = m_pSource->CreatePresentationDescriptor(&pPD);
    if (FAILED(hr))
    {
        goto done;
    }

    BOOL fSelected;
    hr = pPD->GetStreamDescriptorByIndex(0, &fSelected, &pSD);
    if (FAILED(hr))
    {
        goto done;
    }

    hr = pSD->GetMediaTypeHandler(&pHandler);
    if (FAILED(hr))
    {
        goto done;
    }

    DWORD cTypes = 0;
    hr = pHandler->GetMediaTypeCount(&cTypes);
    if (FAILED(hr))
    {
        goto done;
    }
	*pcCount = cTypes;
	*pppTypes = (IMFMediaType**)CoTaskMemAlloc( cTypes * sizeof(IMFMediaType*));

    for (DWORD i = 0; i < cTypes; i++)
    {
        hr = pHandler->GetMediaTypeByIndex(i, &pType);
        if (FAILED(hr))
        {
            goto done;
        }
		(*pppTypes)[i] = pType;
		pType->AddRef();
		// LogMediaType(pType);
        //OutputDebugString(L"\n");

        SafeRelease(&pType);
    }

done:
    SafeRelease(&pPD);
    SafeRelease(&pSD);
    SafeRelease(&pHandler);
    SafeRelease(&pType);
    return hr;
}

HRESULT CPreview::SetFormat(IMFMediaType *pType)
{
    IMFPresentationDescriptor *pPD = NULL;
    IMFStreamDescriptor *pSD = NULL;
    IMFMediaTypeHandler *pHandler = NULL;
	UINT32 width;
	UINT32 height;
	UINT32 sample_size;
    //IMFMediaType *pType = NULL;
	EnterCriticalSection(&m_critsec);

    HRESULT hr = m_pSource->CreatePresentationDescriptor(&pPD);
    if (FAILED(hr))
    {
        goto done;
    }

    BOOL fSelected;
    hr = pPD->GetStreamDescriptorByIndex(0, &fSelected, &pSD);
    if (FAILED(hr))
    {
        goto done;
    }

    hr = pSD->GetMediaTypeHandler(&pHandler);
    if (FAILED(hr))
    {
        goto done;
    }

    /*hr = pHandler->GetMediaTypeByIndex(dwFormatIndex, &pType);
    if (FAILED(hr))
    {
        goto done;
    }*/

    hr = pHandler->SetCurrentMediaType(pType);
	if (FAILED(hr))
    {
        goto done;
    }
	LeaveCriticalSection(&m_critsec);

	hr = TryMediaType(pType);

	EnterCriticalSection(&m_critsec);

	if (FAILED(hr))
    {
        goto done;
    }
	hr = MFGetAttributeSize( pType, MF_MT_FRAME_SIZE, &width, &height);
	if (FAILED(hr))
    {
        goto done;
    }
	hr = pType->GetUINT32( MF_MT_SAMPLE_SIZE, &sample_size);
	if (FAILED(hr))
    {
        goto done;
    }
	LogMediaType(pType);
	
		// Initialize AR Anayse
	if ( g_pAnalyse )
		g_pAnalyse->InitialiseTracker( width, height, sample_size );
	
done:
    SafeRelease(&pPD);
    SafeRelease(&pSD);
    SafeRelease(&pHandler);
    //SafeRelease(&pType);

	LeaveCriticalSection(&m_critsec);

    return hr;
}