#include <iostream>
#include <windows.h>

using namespace std;

bool loadData(const wchar_t *srcFile, void ** ppDataBuffer, size_t * pcbDataBuffer) {
   HANDLE handle = NULL;
   bool result = false;

   if (wcscmp(L"-", srcFile) == 0) {
      handle = GetStdHandle(STD_INPUT_HANDLE);
   }
   else {
      handle = CreateFileW(srcFile, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_ALWAYS, 0, NULL);
   }

   if (!handle) {
      wcout << L"Error opening source file \"" << srcFile << L"\": " << GetLastError() << endl;
   }
   else {
      *pcbDataBuffer = 4000000;
      *ppDataBuffer = (void*)new char[*pcbDataBuffer];

      DWORD bytesRead = *pcbDataBuffer;

      BOOL readResult = ReadFile(handle, *ppDataBuffer, *pcbDataBuffer, &bytesRead, NULL);

      if (!readResult) {
         wcout << L"Error reading data from source file \"" << srcFile << L"\": " << GetLastError() << endl;
      }
      else if (bytesRead == *pcbDataBuffer) {
         wcout << L"Error source file \"" << srcFile << L"\" too big.  Only files less than " << bytesRead << L" bytes supported." << endl;
      }
      else {
         *pcbDataBuffer = bytesRead;
         result = true;
      }
   }
   return result;
}

void modifyResource(const wchar_t *dllPath, const wchar_t *srcFile, BOOL clearExisting, LPCWSTR typeName, LPCWSTR resourceName, WORD lang) {
   HANDLE hDll = BeginUpdateResourceW(dllPath, clearExisting);
   if (NULL == hDll) {
      wcout << L"BeginUpdateResource failed on \"" << dllPath << L"\": " << GetLastError() << endl;
   }
   else {
      void * pDataBuffer = NULL;
      size_t cbDataBuffer = 0;

      if (loadData(srcFile, &pDataBuffer, &cbDataBuffer)) {
         UpdateResourceW(hDll, typeName, resourceName, lang, pDataBuffer, cbDataBuffer);
         EndUpdateResource(hDll, FALSE);
         delete [] pDataBuffer;
      }
      else
         EndUpdateResource(hDll, TRUE);
   }
}

void wmain(int argc, const wchar_t * argv[]) {
   BOOL clearExistingResources = false;
   const wchar_t * dllName = NULL;
   const wchar_t * srcFile = L"-";
   WORD language = 0x0;
   const wchar_t * typeName = NULL;
   const wchar_t * resourceName = NULL;
   bool fail = argc < 2;
   const wchar_t * failedOn = NULL;

   for (int idx = 1; !fail && idx < argc; ++idx) {
      if (wcscmp(argv[idx], L"--dest") == 0) {
         dllName = argv[++idx];
      }
      else if (wcscmp(argv[idx], L"--src") == 0) {
         srcFile = argv[++idx];
      }
      else if (wcscmp(argv[idx], L"--language") == 0) {
         language = _wtoi(argv[++idx]);
      }
      else if (wcscmp(argv[idx], L"--type") == 0) {
         typeName = argv[++idx];
      }
      else if (wcscmp(argv[idx], L"--typei") == 0) {
         typeName = (const wchar_t*)MAKEINTRESOURCE(_wtoi(argv[++idx]));
      }
      else if (wcscmp(argv[idx], L"--name") == 0) {
         resourceName = argv[++idx];
      }
      else if (wcscmp(argv[idx], L"--namei") == 0) {
         resourceName = (const wchar_t*)MAKEINTRESOURCE(_wtoi(argv[++idx]));
      }
      else if (wcscmp(argv[idx], L"--clear") == 0) {
         clearExistingResources = TRUE;
      }
      else {
         failedOn = argv[idx];
         fail = true;
      }
   }

   if (!fail) {
      fail = true;

      if (!dllName)
         wcout << "No DLL specified for modification." << endl;
      else if (!typeName)
         wcout << "No type name specified." << endl;
      else if (!resourceName)
         wcout << "No resource name specified." << endl;
      else {
         modifyResource(dllName, srcFile, clearExistingResources, typeName, resourceName, language);
         fail = false;
      }
   }

   if (fail) {
      if (failedOn)
         wcout << L"Unknown argument: " << failedOn << endl;
      wcout << L"Usage:" << endl << L"\t" << argv[0] << " --dest [dll path] --src [src file] --language [lang id] --type [type name] --name [res name]" << endl;
   }
}