#include "IconChanger.h"
#include <iostream>

using namespace std;
#include <QMessageBox>

QString ICONPATH;
QString EXEPATH;       /* Exe path with one \  */
QString EXEPATHPADDED; /* Exe path with two \\ */


BOOL ResLang( HMODULE hModule, LPCTSTR lpszType, LPCTSTR lpszName, WORD wIDLanguage, LONG_PTR lParam)
{
    if (!IS_INTRESOURCE(lpszName))
    {
//        cout << "LANG -: " << wIDLanguage << endl;
//        cout << "TYPE -: " << lpszType << endl;
//        cout << "NAME -: " << lpszName << endl;


        FreeLibrary( hModule );

        HANDLE hCleanUpdate = BeginUpdateResource( EXEPATHPADDED.toStdString().c_str() , FALSE );
        if ( hCleanUpdate == NULL )
        {
            QMessageBox mess;
            mess.setText(QString("Start Clean UP FAILED: ").append(QString::number( GetLastError() ) ) );
            mess.exec();
        }

        if ( !UpdateResource( hCleanUpdate, lpszType, lpszName, wIDLanguage, NULL, 0) )
        {
            QMessageBox mess;
            mess.setText(QString("Clean UP FAILED: : ").append(QString::number( GetLastError() ) ) );
            mess.exec();
            return FALSE;
        }


        if ( !EndUpdateResource( hCleanUpdate, FALSE) )
        {
            QMessageBox mess;
            mess.setText(QString(" End Clean UP FAILED: ").append(QString::number( GetLastError() ) ) );
            mess.exec();
        }

        hModule = LoadLibrary( EXEPATH.toStdString().c_str() );
        if ( hModule == NULL )
        {
            QMessageBox mess;
            mess.setText(QString("Unable to get handle to EXE/DLL\n").append(EXEPATH).append("\nError Code: ").append( QString::number(GetLastError()) ) );
            mess.exec();
        }

    } else {
//        cout << "LANG -: " << (long)wIDLanguage << endl;
//        cout << "TYPE -: " << (long)lpszType << endl;
//        cout << "NAME -: " << (long)lpszName << endl;
//        cout << "FILE -: " << EXEPATHPADDED.toStdString().c_str() << endl;

        FreeLibrary( hModule );

        HANDLE hCleanUpdate = BeginUpdateResource( EXEPATHPADDED.toStdString().c_str() , FALSE );
        if ( hCleanUpdate == NULL )
        {
            QMessageBox mess;
            mess.setText(QString("Start Clean UP FAILED: ").append(QString::number( GetLastError() ) ) );
            mess.exec();
        }

        if ( !UpdateResource( hCleanUpdate, MAKEINTRESOURCE((long)lpszType), MAKEINTRESOURCE((long)lpszName), wIDLanguage, NULL, 0) )
        {
            QMessageBox mess;
            mess.setText(QString("Clean UP FAILED: : ").append(QString::number( GetLastError() ) ) );
            mess.exec();
            return FALSE;
        }


        if ( !EndUpdateResource( hCleanUpdate, FALSE) )
        {
            QMessageBox mess;
            mess.setText(QString(" End Clean UP FAILED: ").append(QString::number( GetLastError() ) ) );
            mess.exec();
        }

        hModule = LoadLibrary( EXEPATH.toStdString().c_str() );
        if ( hModule == NULL )
        {
            QMessageBox mess;
            mess.setText(QString("Unable to get handle to EXE/DLL\n").append(EXEPATH).append("\nError Code: ").append( QString::number(GetLastError()) ) );
            mess.exec();
        }

    }

    return TRUE;
}

BOOL ResNames(HMODULE hModule, LPCTSTR lpType, LPTSTR lpName, LONG lParam)
{
    if (!IS_INTRESOURCE(lpName))
    {
//        cout << "is int resource. NAMES" << endl;
//        cout << "TYPE in Names: " << lpType << endl;
//        cout << "NAME in Names: " << lpName << endl;
        EnumResourceLanguages( hModule, lpType, lpName, (ENUMRESLANGPROC)ResLang, 0);

    } else {
//        cout << "is int resource. NAMES" << endl;
//        cout << "TYPE in Names: " << (long)lpType << endl;
//        cout << "NAME in Names: " << (long)lpName << endl;
        EnumResourceLanguages( hModule, lpType, lpName, (ENUMRESLANGPROC)ResLang, 0);
    }

   return TRUE;
}

BOOL ResTypes(HMODULE hModule, LPTSTR lpType, LONG lParam)
{
    if (!IS_INTRESOURCE(lpType))
    {
//      cout << "INT is not resource TYPE" << endl;
//      cout << "TYPE:" << lpType << endl;
        if ( (lpType == RT_ICON ) || (lpType == (RT_ICON + 11)) )
        {
            EnumResourceNames(hModule, lpType, (ENUMRESNAMEPROC)ResNames, 0);
        }


    } else {
//        cout << "INT is resource TYPE" << endl;
//        cout << "TYPE:" << (long)lpType << endl;
        if ( (lpType == RT_ICON ) || (lpType == (RT_ICON + 11)) )
        {
            EnumResourceNames(hModule, lpType, (ENUMRESNAMEPROC)ResNames, 0);
        }
    }

   return TRUE;
}


IconChanger::IconChanger() // Constructor
{
    hExeSource = NULL;
    hrResSource = NULL;
    hGlobSource = NULL;
    lpIconBuffer = NULL;
    nIconCount = 0;
}
                                                            // Retrive Icon Buffer
void IconChanger::GetIconFromSource(QString sIconPath, QString sExePath) // from source Exe/DLL
{
    ICONPATH = sIconPath; EXEPATH = sExePath;

    LPICONDIRENTRY pIconDirEntry(NULL);
    LPGRPICONDIR pGrpIconDir(NULL);
    HEADER header;
    LPBYTE pIconBytes(NULL);
    HANDLE hIconFile(NULL);
    DWORD dwRet(0), nSize(0), nGSize(0), dwReserved(0);
    HANDLE hUpdate(NULL);
    BOOL ret(FALSE);
    WORD i(0);
    //sIconPath.replace("\\", "\\\\", Qt::CaseSensitive);
    //sExePath.replace("\\", "\\\\", Qt::CaseSensitive);

    hIconFile = CreateFile( sIconPath.toStdString().c_str(), GENERIC_READ,
                            0, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if ( hIconFile == INVALID_HANDLE_VALUE )
    {
        QMessageBox mess; mess.setText("Unable to open file.\n" + sIconPath );
        mess.exec();
    }

    ret = ReadFile( hIconFile, &header, sizeof(HEADER), &dwReserved, NULL );
    if (!ret)
    {
        CloseHandle(hExeSource);
        QMessageBox mess; mess.setText("Unable to Read Icon File " + sIconPath );
        mess.exec();
        return;
    }

    pIconDirEntry = (LPICONDIRENTRY)new BYTE[ header.idCount * sizeof(ICONDIRENTRY) ];
    if ( pIconDirEntry == NULL )
    {
        CloseHandle(hIconFile);
        QMessageBox mess; mess.setText("Unable to allocate memory for Icon Group Header" + sIconPath );
        mess.exec();
        return;
    }

    ret = ReadFile(hIconFile, pIconDirEntry, header.idCount*sizeof(ICONDIRENTRY), &dwReserved, NULL);
    if (!ret)
    {
        delete[] pIconDirEntry;
        CloseHandle(hIconFile);
        QMessageBox mess;
        mess.setText("Unable to read Icon Dir Entrys" + sIconPath );
        mess.exec();
        return;
    }

    nGSize = sizeof(GRPICONDIR)+header.idCount*sizeof(ICONDIRENTRY);
    pGrpIconDir = (LPGRPICONDIR)new BYTE[nGSize];
    ZeroMemory( pGrpIconDir, nSize );

    pGrpIconDir->idReserved = header.idReserved;
    pGrpIconDir->idType = header.idType;
    pGrpIconDir->idCount = header.idCount;

   // nIconCount = 0;

    for ( i = 0; i < header.idCount; i++ )
    {
        pGrpIconDir->idEntries[i].bWidth = pIconDirEntry[i].bWidth;
        pGrpIconDir->idEntries[i].bHeight = pIconDirEntry[i].bHeight;
        pGrpIconDir->idEntries[i].bColorCount = pIconDirEntry[i].bColorCount;
        pGrpIconDir->idEntries[i].bReserved = pIconDirEntry[i].bReserved;
        pGrpIconDir->idEntries[i].wPlanes = pIconDirEntry[i].wPlanes;
        pGrpIconDir->idEntries[i].wBitCount = pIconDirEntry[i].wBitCount;
        pGrpIconDir->idEntries[i].dwBytesInRes = pIconDirEntry[i].dwBytesInRes;
        pGrpIconDir->idEntries[i].nID = i+1;
    }

    EXEPATHPADDED = EXEPATH;
    EXEPATHPADDED.replace("\\", "\\\\", Qt::CaseInsensitive);
    hExeSource = LoadLibrary( sExePath.toStdString().c_str() );
    if ( hExeSource == NULL )
    {
        QMessageBox mess;
        mess.setText("Unable to get handle to EXE/DLL\n" + sExePath );
        mess.exec();
    }

    BOOL more = true;
    while(more){
        more = EnumResourceTypes( hExeSource, (ENUMRESTYPEPROC)ResTypes, 0);
    }

    CloseHandle(hIconFile);
    FreeLibrary( hExeSource );
    QMessageBox mess;
    mess.setText("Cleanup DONE!... " );
    mess.exec();
}

IconChanger::~IconChanger() // De-Constructor
{

}