#include <windows.h>
#include <tchar.h>
#include <toolhelp.h>

#include "resource.h"

#if (_WIN32_WCE < 200)
#error "Must be compiled for Windows CE 2.0 or later."
#endif


// stolen from rapi.h - not included in WCE headers
typedef enum tagRAPISTREAMFLAG
{
	STREAM_TIMEOUT_READ
} RAPISTREAMFLAG;

DECLARE_INTERFACE_ (IRAPIStream,  IStream)
{
    STDMETHOD(SetRapiStat)( THIS_ RAPISTREAMFLAG Flag, DWORD dwValue) PURE;
    STDMETHOD(GetRapiStat)( THIS_ RAPISTREAMFLAG Flag, DWORD *pdwValue) PURE;
};


HINSTANCE g_hInstance = NULL;

// max size for text box in dialog
const int MAX_CHARS = 32;
const int TEXT_MAX = MAX_CHARS+1;   


STDAPI_(BOOL) APIENTRY DllMain(HANDLE hDll, DWORD dwReason, LPVOID lpReserved)
{
    switch( dwReason )
    {
        case DLL_PROCESS_ATTACH:
            g_hInstance = hDll;
            break;

        case DLL_PROCESS_DETACH:
        case DLL_THREAD_ATTACH:
        case DLL_THREAD_DETACH:
            break;
    }

    return TRUE;
}


STDAPI GetLoadedModules(
   DWORD cbInput,
   BYTE  *pInput,
   DWORD *pcbOutput,
   BYTE  **ppOutput,
   IRAPIStream *pIRAPIStream )
{


    HANDLE          hSnap;
    MODULEENTRY32   me = { sizeof(MODULEENTRY32) };
    LPWSTR          lpwszMods = NULL;
    DWORD           dwCurSize = 0;


    // init output variables
    *ppOutput = NULL;
    *pcbOutput = 0;

    // get module list
    hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE | TH32CS_GETALLMODS, 0);
    if ( hSnap == (HANDLE)-1 )
    {
        return S_FALSE;
    }

    BOOL bFound = Module32First(hSnap, &me);

    while ( bFound == TRUE )
    {

        // package up modules as CRLF delimited list
        dwCurSize += (wcslen(me.szModule)*sizeof(TCHAR)) + (sizeof(TCHAR)*2);

        if ( lpwszMods == NULL )
        lpwszMods = (LPWSTR)LocalAlloc(LPTR, dwCurSize);
        else
         lpwszMods = (LPWSTR)LocalReAlloc((HLOCAL)lpwszMods, dwCurSize, LMEM_MOVEABLE);

        if ( lpwszMods == NULL )
        {
         dwCurSize = 0;
         break;
        }

        wcscat(lpwszMods, me.szModule);
        wcscat(lpwszMods, L"\n");

        bFound = Module32Next(hSnap, &me);
    }

    CloseToolhelp32Snapshot(hSnap);


    // fill in output vars
    *ppOutput = (PBYTE)lpwszMods;
    *pcbOutput = dwCurSize;

    return S_OK;
}


BOOL CALLBACK DlgProc(HWND hwndDlg, 
                      UINT uMsg, 
                      WPARAM wParam, 
                      LPARAM lParam); 



STDAPI DlgStreamTest(
   DWORD cbInput,
   BYTE  *pInput,
   DWORD *pcbOutput,
   BYTE  **ppOutput,
   IRAPIStream *pIRAPIStream )
{

    // init output variables
    *pcbOutput = 0;
    *ppOutput = NULL;

    // display the dialog, passing the stream in lParam
    DialogBoxParam(g_hInstance,
                   MAKEINTRESOURCE(IDD_MAINDLG), 
                   NULL, 
                   (DLGPROC)DlgProc,
                   (LPARAM)pIRAPIStream);

    return S_OK;
}

// simple struct to pass to thread function
struct ReceiveData
{
    HWND hwnd;
    IRAPIStream *pStream;
};


DWORD WINAPI ReceiveThread(PVOID pv)
{
    DWORD       cbOut, cbRead;
    HRESULT     hr;
    ReceiveData  *prd = (ReceiveData*)pv;
    TCHAR       chBuf[TEXT_MAX];

    if ( prd == NULL || prd->pStream == NULL || !IsWindow(prd->hwnd) )
    {
        OutputDebugString(L"ERROR: invalid parameter to ReceiveThread\n");

        // release the stream if we can
        if ( prd && prd->pStream )
            prd->pStream->Release();

        return 0;
    }

    while ( TRUE )
    {
        hr = prd->pStream->Read( &cbOut, sizeof ( cbOut ), &cbRead );

        if ( SUCCEEDED(hr) )
        {
            hr = prd->pStream->Read( (PBYTE)chBuf, cbOut , &cbRead );

            if ( SUCCEEDED(hr) )
            {
                if ( wcscmp(chBuf, L"\x01|done|\x02") == 0 )
                    break;

                SetWindowText(GetDlgItem(prd->hwnd, IDC_TEXT), chBuf);
            }
            else
            {
                OutputDebugString(_T("ERROR: read size failed...\n"));
                break;
            }
        }
        else
        {
            OutputDebugString(_T("ERROR: read data failed...\n"));
            break;
        }
    }

    // Post a message back to the dialog to let it know we are finished
    // DO NOT change this to send because the dialog waits for the thread to exit
    PostMessage(prd->hwnd, WM_COMMAND, MAKELONG(IDC_FINISHED,BN_CLICKED), NULL);

    // release the stream interface
    prd->pStream->Release();

    return 0;
}


BOOL CALLBACK DlgProc(HWND hwndDlg, 
                      UINT uMsg, 
                      WPARAM wParam, 
                      LPARAM lParam)
{

    BOOL    bRet = FALSE;

    static ReceiveData rd;
    static HANDLE hReceiveThread = NULL;


    switch( uMsg )
    {
        case WM_INITDIALOG:
            {
                // limit the text box to MAX_CHARS
                SendMessage(GetDlgItem(hwndDlg, IDC_TEXT), EM_LIMITTEXT, MAX_CHARS, 0);

                // start the receive thread
                rd.hwnd = hwndDlg;
                rd.pStream = (IRAPIStream*)lParam;

                DWORD dwTID;
                hReceiveThread = CreateThread(NULL, 0, ReceiveThread, &rd, 0, &dwTID);
                if ( hReceiveThread == NULL )
                {
                    OutputDebugString(_T("ERROR: could not create ReceiveThread...\n"));
                    rd.pStream->Release();
                }
                else
                {
                    // make sure the window is visible
                    SetForegroundWindow(hwndDlg);
                }
                bRet = TRUE;
            }
            break;

        case WM_COMMAND:
        {
            int     nID = LOWORD(wParam);
            int     wNotify = HIWORD(wParam);

            // end the dialog if we received an IDC_FINISHED notification
            if ( wNotify == BN_CLICKED && nID == IDC_FINISHED )
            {
                // wait for ReceiveThread to exit and kill the dialog

                if ( hReceiveThread )
                {
                    WaitForSingleObject(hReceiveThread, INFINITE);
                    hReceiveThread = NULL;
                }

                EndDialog(hwndDlg, IDC_FINISHED);
                bRet = TRUE;
            }
        }
        break;
    }

    return bRet;
}

