nfd2 0.3.1

/*
  Native File Dialog

  http://www.frogtoss.com/labs
 */


#ifdef __MINGW32__
// Explicitly setting NTDDI version, this is necessary for the MinGW compiler
#    define NTDDI_VERSION NTDDI_VISTA
#    define _WIN32_WINNT _WIN32_WINNT_VISTA
#endif

#ifndef _CRT_SECURE_NO_WARNINGS
#    define _CRT_SECURE_NO_WARNINGS
#endif

#define _CRTDBG_MAP_ALLOC
#include <stdlib.h>
#include <crtdbg.h>

/* only locally define UNICODE in this compilation unit */
#ifndef UNICODE
#    define UNICODE
#endif

#include <wchar.h>
#include <stdio.h>
#include <assert.h>
#include <windows.h>
#include <shobjidl.h>
#include "nfd_common.h"

#define COM_INITFLAGS ::COINIT_APARTMENTTHREADED | ::COINIT_DISABLE_OLE1DDE

static BOOL
COMIsInitialized(HRESULT coResult)
{
    if (coResult == RPC_E_CHANGED_MODE) {
        // If COM was previously initialized with different init flags,
        // NFD still needs to operate. Eat this warning.
        return TRUE;
    }

    return SUCCEEDED(coResult);
}

static HRESULT
COMInit(void)
{
    return ::CoInitializeEx(NULL, COM_INITFLAGS);
}

static void
COMUninit(HRESULT coResult)
{
    // do not uninitialize if RPC_E_CHANGED_MODE occurred -- this
    // case does not refcount COM.
    if (SUCCEEDED(coResult))
        ::CoUninitialize();
}

// allocs the space in outPath -- call free()
static void
CopyWCharToNFDChar(const wchar_t* inStr, nfdchar_t** outStr)
{
    int inStrCharacterCount = static_cast<int>(wcslen(inStr));
    int bytesNeeded = WideCharToMultiByte(
        CP_UTF8, 0, inStr, inStrCharacterCount, NULL, 0, NULL, NULL);
    assert(bytesNeeded);
    bytesNeeded += 1;

    *outStr = (nfdchar_t*)NFDi_Malloc(bytesNeeded);
    if (!*outStr)
        return;

    int bytesWritten =
        WideCharToMultiByte(CP_UTF8, 0, inStr, -1, *outStr, bytesNeeded, NULL, NULL);
    assert(bytesWritten);
    _NFD_UNUSED(bytesWritten);
}

/* includes NULL terminator byte in return */
static size_t
GetUTF8ByteCountForWChar(const wchar_t* str)
{
    size_t bytesNeeded = WideCharToMultiByte(CP_UTF8, 0, str, -1, NULL, 0, NULL, NULL);
    assert(bytesNeeded);
    return bytesNeeded + 1;
}

// write to outPtr -- no free() necessary.
static int
CopyWCharToExistingNFDCharBuffer(const wchar_t* inStr, nfdchar_t* outPtr)
{
    int bytesNeeded = static_cast<int>(GetUTF8ByteCountForWChar(inStr));

    /* invocation copies null term */
    int bytesWritten =
        WideCharToMultiByte(CP_UTF8, 0, inStr, -1, outPtr, bytesNeeded, NULL, 0);
    assert(bytesWritten);

    return bytesWritten;
}


// allocs the space in outStr -- call free()
static void
CopyNFDCharToWChar(const nfdchar_t* inStr, size_t inStrLen, wchar_t** outStr)
{
    int inStrByteCount = (int)inStrLen; /* byte count does not include terminator */
    int charsNeeded =
        MultiByteToWideChar(CP_UTF8, 0, inStr, (int)inStrByteCount, NULL, 0);
    assert(charsNeeded);
    assert(!*outStr);
    charsNeeded += 1;  // terminator

    *outStr = (wchar_t*)NFDi_Malloc(charsNeeded * sizeof(wchar_t));
    if (!*outStr)
        return;

    MultiByteToWideChar(CP_UTF8, 0, inStr, (int)inStrByteCount, *outStr, charsNeeded);
    (*outStr)[charsNeeded - 1] = '\0';
}


/* ext is in format "jpg", no wildcards or separators */
static int
AppendExtensionToSpecBuf(const char* ext, char* specBuf, size_t specBufLen)
{
    const char SEP[] = ";";
    assert(specBufLen > strlen(ext) + 3);

    if (strlen(specBuf) > 0) {
        strncat(specBuf, SEP, specBufLen - strlen(specBuf) - 1);
        specBufLen += strlen(SEP);
    }

    char extWildcard[NFD_MAX_STRLEN];
    int  bytesWritten = sprintf_s(extWildcard, NFD_MAX_STRLEN, "*.%s", ext);
    assert(bytesWritten == (int)(strlen(ext) + 2));
    _NFD_UNUSED(bytesWritten);

    strncat(specBuf, extWildcard, specBufLen - strlen(specBuf) - 1);

    return NFD_OKAY;
}

static nfdresult_t
AddFiltersToDialog(::IFileDialog* fileDialog, const char* filterList)
{
    const wchar_t WILDCARD[] = L"*.*";

    if (!filterList || strlen(filterList) == 0)
        return NFD_OKAY;

    // Count rows to alloc
    UINT filterCount =
        1; /* guaranteed to have one filter on a correct, non-empty parse */
    const char* p_filterList;
    for (p_filterList = filterList; *p_filterList; ++p_filterList) {
        if (*p_filterList == ';')
            ++filterCount;
    }

    assert(filterCount);
    if (!filterCount) {
        NFDi_SetError("Error parsing filters.");
        return NFD_ERROR;
    }

    /* filterCount plus 1 because we hardcode the *.* wildcard after the while loop */
    COMDLG_FILTERSPEC* specList = (COMDLG_FILTERSPEC*)NFDi_Malloc(
        sizeof(COMDLG_FILTERSPEC) * ((size_t)filterCount + 1));
    if (!specList) {
        return NFD_ERROR;
    }
    for (UINT i = 0; i < filterCount + 1; ++i) {
        specList[i].pszName = NULL;
        specList[i].pszSpec = NULL;
    }

    size_t specIdx = 0;
    p_filterList = filterList;
    char  typebuf[NFD_MAX_STRLEN] = {0}; /* one per comma or semicolon */
    char* p_typebuf = typebuf;

    char specbuf[NFD_MAX_STRLEN] = {0}; /* one per semicolon */

    while (1) {
        if (NFDi_IsFilterSegmentChar(*p_filterList)) {
            if (specbuf[0] == '\0') {
                wchar_t* ext = NULL;
                CopyNFDCharToWChar(typebuf, strlen(typebuf), &ext);
                fileDialog->SetDefaultExtension(ext);
                NFDi_Free(ext);
            }

            /* append a type to the specbuf (pending filter) */
            AppendExtensionToSpecBuf(typebuf, specbuf, NFD_MAX_STRLEN);

            p_typebuf = typebuf;
            memset(typebuf, 0, sizeof(char) * NFD_MAX_STRLEN);
        }

        if (*p_filterList == ';' || *p_filterList == '\0') {
            /* end of filter -- add it to specList */

            CopyNFDCharToWChar(
                specbuf, strlen(specbuf), (wchar_t**)&specList[specIdx].pszName);
            CopyNFDCharToWChar(
                specbuf, strlen(specbuf), (wchar_t**)&specList[specIdx].pszSpec);

            memset(specbuf, 0, sizeof(char) * NFD_MAX_STRLEN);
            ++specIdx;
            if (specIdx == filterCount)
                break;
        }

        if (!NFDi_IsFilterSegmentChar(*p_filterList)) {
            *p_typebuf = *p_filterList;
            ++p_typebuf;
        }

        ++p_filterList;
    }

    /* Add wildcard */
    specList[specIdx].pszSpec = WILDCARD;
    specList[specIdx].pszName = WILDCARD;

    fileDialog->SetFileTypes(filterCount + 1, specList);

    /* free speclist */
    for (size_t i = 0; i < filterCount; ++i) {
        NFDi_Free((void*)specList[i].pszSpec);
    }
    NFDi_Free(specList);

    return NFD_OKAY;
}

static nfdresult_t
AllocPathSet(IShellItemArray* shellItems, nfdpathset_t* pathSet)
{
    const char ERRORMSG[] = "Error allocating pathset.";

    assert(shellItems);
    assert(pathSet);

    // How many items in shellItems?
    DWORD   numShellItems;
    HRESULT result = shellItems->GetCount(&numShellItems);
    if (!SUCCEEDED(result)) {
        NFDi_SetError(ERRORMSG);
        return NFD_ERROR;
    }

    pathSet->count = static_cast<size_t>(numShellItems);
    assert(pathSet->count > 0);

    pathSet->indices = (size_t*)NFDi_Malloc(sizeof(size_t) * pathSet->count);
    if (!pathSet->indices) {
        return NFD_ERROR;
    }

    /* count the total bytes needed for buf */
    size_t bufSize = 0;
    for (DWORD i = 0; i < numShellItems; ++i) {
        ::IShellItem* shellItem;
        result = shellItems->GetItemAt(i, &shellItem);
        if (!SUCCEEDED(result)) {
            NFDi_SetError(ERRORMSG);
            return NFD_ERROR;
        }

        // Confirm SFGAO_FILESYSTEM is true for this shellitem, or ignore it.
        SFGAOF attribs;
        result = shellItem->GetAttributes(SFGAO_FILESYSTEM, &attribs);
        if (!SUCCEEDED(result)) {
            NFDi_SetError(ERRORMSG);
            return NFD_ERROR;
        }
        if (!(attribs & SFGAO_FILESYSTEM))
            continue;

        LPWSTR name;
        shellItem->GetDisplayName(SIGDN_FILESYSPATH, &name);

        // Calculate length of name with UTF-8 encoding
        bufSize += GetUTF8ByteCountForWChar(name);

        CoTaskMemFree(name);
    }

    assert(bufSize);

    pathSet->buf = (nfdchar_t*)NFDi_Malloc(sizeof(nfdchar_t) * bufSize);
    memset(pathSet->buf, 0, sizeof(nfdchar_t) * bufSize);

    /* fill buf */
    nfdchar_t* p_buf = pathSet->buf;
    for (DWORD i = 0; i < numShellItems; ++i) {
        ::IShellItem* shellItem;
        result = shellItems->GetItemAt(i, &shellItem);
        if (!SUCCEEDED(result)) {
            NFDi_SetError(ERRORMSG);
            return NFD_ERROR;
        }

        // Confirm SFGAO_FILESYSTEM is true for this shellitem, or ignore it.
        SFGAOF attribs;
        result = shellItem->GetAttributes(SFGAO_FILESYSTEM, &attribs);
        if (!SUCCEEDED(result)) {
            NFDi_SetError(ERRORMSG);
            return NFD_ERROR;
        }
        if (!(attribs & SFGAO_FILESYSTEM))
            continue;

        LPWSTR name;
        shellItem->GetDisplayName(SIGDN_FILESYSPATH, &name);

        int bytesWritten = CopyWCharToExistingNFDCharBuffer(name, p_buf);
        CoTaskMemFree(name);

        ptrdiff_t index = p_buf - pathSet->buf;
        assert(index >= 0);
        pathSet->indices[i] = static_cast<size_t>(index);

        p_buf += bytesWritten;
    }

    return NFD_OKAY;
}


static nfdresult_t
SetDefaultDir(IFileDialog* dialog, const char* defaultPath)
{
    if (!defaultPath || defaultPath[0] == '\0')
        return NFD_OKAY;

    const char *defaultDir, *defaultFilename;
    NFDi_SplitPath(defaultPath, &defaultDir, &defaultFilename);

    size_t defaultDirLen;
    if (defaultFilename) {
        assert(defaultFilename > defaultDir);
        defaultDirLen = defaultFilename - defaultDir;
    } else {
        defaultDirLen = strlen(defaultDir);
    }

    wchar_t* defaultPathW = {0};
    CopyNFDCharToWChar(defaultPath, defaultDirLen, &defaultPathW);

    IShellItem* folder;
    HRESULT     result =
        SHCreateItemFromParsingName(defaultPathW, NULL, IID_PPV_ARGS(&folder));

    // Valid non results.
    if (result == HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND) ||
        result == HRESULT_FROM_WIN32(ERROR_INVALID_DRIVE)) {
        NFDi_Free(defaultPathW);
        return NFD_OKAY;
    }

    if (!SUCCEEDED(result)) {
        NFDi_SetError("Error creating ShellItem");
        NFDi_Free(defaultPathW);
        return NFD_ERROR;
    }

    // Could also call SetDefaultFolder(), but this guarantees defaultPath -- more consistency across API.
    dialog->SetFolder(folder);

    NFDi_Free(defaultPathW);
    folder->Release();

    return NFD_OKAY;
}

static void
SetDefaultFile(IFileDialog* dialog, const char* defaultPath)
{
    if (!defaultPath || defaultPath[0] == '\0')
        return;

    const char *defaultDir, *defaultFilename;
    NFDi_SplitPath(defaultPath, &defaultDir, &defaultFilename);

    // no filename in to set -- nothing to do
    if (!defaultFilename)
        return;

    wchar_t* defaultFilenameW = {0};
    CopyNFDCharToWChar(defaultFilename, strlen(defaultFilename), &defaultFilenameW);

    dialog->SetFileName(defaultFilenameW);

    NFDi_Free(defaultFilenameW);
}

/* public */


nfdresult_t
NFD_OpenDialog(const nfdchar_t* filterList, const nfdchar_t* defaultPath, nfdchar_t** outPath)
{
    nfdresult_t nfdResult = NFD_ERROR;


    HRESULT coResult = COMInit();
    if (!COMIsInitialized(coResult)) {
        NFDi_SetError("Could not initialize COM.");
        return nfdResult;
    }

    // Create dialog
    ::IFileOpenDialog* fileOpenDialog(NULL);
    HRESULT            result = ::CoCreateInstance(::CLSID_FileOpenDialog,
                                        NULL,
                                        CLSCTX_ALL,
                                        ::IID_IFileOpenDialog,
                                        reinterpret_cast<void**>(&fileOpenDialog));

    if (!SUCCEEDED(result)) {
        NFDi_SetError("Could not create dialog.");
        goto end;
    }

    // Build the filter list
    if (!AddFiltersToDialog(fileOpenDialog, filterList)) {
        goto end;
    }

    // Set the default dir
    if (!SetDefaultDir(fileOpenDialog, defaultPath)) {
        goto end;
    }

    // Set the default filename
    SetDefaultFile(fileOpenDialog, defaultPath);

    // Show the dialog.
    result = fileOpenDialog->Show(NULL);
    if (SUCCEEDED(result)) {
        // Get the file name
        ::IShellItem* shellItem(NULL);
        result = fileOpenDialog->GetResult(&shellItem);
        if (!SUCCEEDED(result)) {
            NFDi_SetError("Could not get shell item from dialog.");
            goto end;
        }
        wchar_t* filePath(NULL);
        result = shellItem->GetDisplayName(::SIGDN_FILESYSPATH, &filePath);
        if (!SUCCEEDED(result)) {
            NFDi_SetError("Could not get file path for selected.");
            shellItem->Release();
            goto end;
        }

        CopyWCharToNFDChar(filePath, outPath);
        CoTaskMemFree(filePath);
        if (!*outPath) {
            /* error is malloc-based, error message would be redundant */
            shellItem->Release();
            goto end;
        }

        nfdResult = NFD_OKAY;
        shellItem->Release();
    } else if (result == HRESULT_FROM_WIN32(ERROR_CANCELLED)) {
        nfdResult = NFD_CANCEL;
    } else {
        NFDi_SetError("File dialog box show failed.");
        nfdResult = NFD_ERROR;
    }

end:
    if (fileOpenDialog)
        fileOpenDialog->Release();

    COMUninit(coResult);

    return nfdResult;
}

nfdresult_t
NFD_OpenDialogMultiple(const nfdchar_t* filterList,
                       const nfdchar_t* defaultPath,
                       nfdpathset_t*    outPaths)
{
    nfdresult_t nfdResult = NFD_ERROR;


    HRESULT coResult = COMInit();
    if (!COMIsInitialized(coResult)) {
        NFDi_SetError("Could not initialize COM.");
        return nfdResult;
    }

    // Create dialog
    ::IFileOpenDialog* fileOpenDialog(NULL);
    HRESULT            result = ::CoCreateInstance(::CLSID_FileOpenDialog,
                                        NULL,
                                        CLSCTX_ALL,
                                        ::IID_IFileOpenDialog,
                                        reinterpret_cast<void**>(&fileOpenDialog));

    if (!SUCCEEDED(result)) {
        fileOpenDialog = NULL;
        NFDi_SetError("Could not create dialog.");
        goto end;
    }

    // Build the filter list
    if (!AddFiltersToDialog(fileOpenDialog, filterList)) {
        goto end;
    }

    // Set the default dir
    if (!SetDefaultDir(fileOpenDialog, defaultPath)) {
        goto end;
    }

    // Set the default filename
    SetDefaultFile(fileOpenDialog, defaultPath);

    // Set a flag for multiple options
    DWORD dwFlags;
    result = fileOpenDialog->GetOptions(&dwFlags);
    if (!SUCCEEDED(result)) {
        NFDi_SetError("Could not get options.");
        goto end;
    }
    result = fileOpenDialog->SetOptions(dwFlags | FOS_ALLOWMULTISELECT);
    if (!SUCCEEDED(result)) {
        NFDi_SetError("Could not set options.");
        goto end;
    }

    // Show the dialog.
    result = fileOpenDialog->Show(NULL);
    if (SUCCEEDED(result)) {
        IShellItemArray* shellItems;
        result = fileOpenDialog->GetResults(&shellItems);
        if (!SUCCEEDED(result)) {
            NFDi_SetError("Could not get shell items.");
            goto end;
        }

        if (AllocPathSet(shellItems, outPaths) == NFD_ERROR) {
            shellItems->Release();
            goto end;
        }

        shellItems->Release();
        nfdResult = NFD_OKAY;
    } else if (result == HRESULT_FROM_WIN32(ERROR_CANCELLED)) {
        nfdResult = NFD_CANCEL;
    } else {
        NFDi_SetError("File dialog box show failed.");
        nfdResult = NFD_ERROR;
    }

end:
    if (fileOpenDialog)
        fileOpenDialog->Release();

    COMUninit(coResult);

    return nfdResult;
}

nfdresult_t
NFD_SaveDialog(const nfdchar_t* filterList, const nfdchar_t* defaultPath, nfdchar_t** outPath)
{
    nfdresult_t nfdResult = NFD_ERROR;

    HRESULT coResult = COMInit();
    if (!COMIsInitialized(coResult)) {
        NFDi_SetError("Could not initialize COM.");
        return nfdResult;
    }

    // Create dialog
    ::IFileSaveDialog* fileSaveDialog(NULL);
    HRESULT            result = ::CoCreateInstance(::CLSID_FileSaveDialog,
                                        NULL,
                                        CLSCTX_ALL,
                                        ::IID_IFileSaveDialog,
                                        reinterpret_cast<void**>(&fileSaveDialog));

    if (!SUCCEEDED(result)) {
        fileSaveDialog = NULL;
        NFDi_SetError("Could not create dialog.");
        goto end;
    }

    // Build the filter list
    if (!AddFiltersToDialog(fileSaveDialog, filterList)) {
        goto end;
    }

    // Set the default dir
    if (!SetDefaultDir(fileSaveDialog, defaultPath)) {
        goto end;
    }

    // Set the default filename
    SetDefaultFile(fileSaveDialog, defaultPath);

    // Show the dialog.
    result = fileSaveDialog->Show(NULL);
    if (SUCCEEDED(result)) {
        // Get the file name
        ::IShellItem* shellItem;
        result = fileSaveDialog->GetResult(&shellItem);
        if (!SUCCEEDED(result)) {
            NFDi_SetError("Could not get shell item from dialog.");
            goto end;
        }
        wchar_t* filePath(NULL);
        result = shellItem->GetDisplayName(::SIGDN_FILESYSPATH, &filePath);
        if (!SUCCEEDED(result)) {
            shellItem->Release();
            NFDi_SetError("Could not get file path for selected.");
            goto end;
        }

        CopyWCharToNFDChar(filePath, outPath);
        CoTaskMemFree(filePath);
        if (!*outPath) {
            /* error is malloc-based, error message would be redundant */
            shellItem->Release();
            goto end;
        }

        nfdResult = NFD_OKAY;
        shellItem->Release();
    } else if (result == HRESULT_FROM_WIN32(ERROR_CANCELLED)) {
        nfdResult = NFD_CANCEL;
    } else {
        NFDi_SetError("File dialog box show failed.");
        nfdResult = NFD_ERROR;
    }

end:
    if (fileSaveDialog)
        fileSaveDialog->Release();

    COMUninit(coResult);

    return nfdResult;
}



nfdresult_t
NFD_PickFolder(const nfdchar_t* defaultPath, nfdchar_t** outPath)
{
    nfdresult_t nfdResult = NFD_ERROR;
    DWORD       dwOptions = 0;

    HRESULT coResult = COMInit();
    if (!COMIsInitialized(coResult)) {
        NFDi_SetError("CoInitializeEx failed.");
        return nfdResult;
    }

    // Create dialog
    ::IFileOpenDialog* fileDialog(NULL);
    HRESULT            result = CoCreateInstance(
        CLSID_FileOpenDialog, NULL, CLSCTX_ALL, IID_PPV_ARGS(&fileDialog));
    if (!SUCCEEDED(result)) {
        NFDi_SetError("CoCreateInstance for CLSID_FileOpenDialog failed.");
        goto end;
    }

    // Set the default dir
    if (SetDefaultDir(fileDialog, defaultPath) != NFD_OKAY) {
        NFDi_SetError("SetDefaultDir failed.");
        goto end;
    }

    // Get the dialogs options
    if (!SUCCEEDED(fileDialog->GetOptions(&dwOptions))) {
        NFDi_SetError("GetOptions for IFileDialog failed.");
        goto end;
    }

    // Add in FOS_PICKFOLDERS which hides files and only allows selection of folders
    if (!SUCCEEDED(fileDialog->SetOptions(dwOptions | FOS_PICKFOLDERS))) {
        NFDi_SetError("SetOptions for IFileDialog failed.");
        goto end;
    }

    // Show the dialog to the user
    result = fileDialog->Show(NULL);
    if (SUCCEEDED(result)) {
        // Get the folder name
        ::IShellItem* shellItem(NULL);

        result = fileDialog->GetResult(&shellItem);
        if (!SUCCEEDED(result)) {
            NFDi_SetError("Could not get file path for selected.");
            shellItem->Release();
            goto end;
        }

        wchar_t* path = NULL;
        result = shellItem->GetDisplayName(SIGDN_DESKTOPABSOLUTEPARSING, &path);
        if (!SUCCEEDED(result)) {
            NFDi_SetError("GetDisplayName for IShellItem failed.");
            shellItem->Release();
            goto end;
        }

        CopyWCharToNFDChar(path, outPath);
        CoTaskMemFree(path);
        if (!*outPath) {
            shellItem->Release();
            goto end;
        }

        nfdResult = NFD_OKAY;
        shellItem->Release();
    } else if (result == HRESULT_FROM_WIN32(ERROR_CANCELLED)) {
        nfdResult = NFD_CANCEL;
    } else {
        NFDi_SetError("Show for IFileDialog failed.");
        nfdResult = NFD_ERROR;
    }

end:

    if (fileDialog)
        fileDialog->Release();

    COMUninit(coResult);

    return nfdResult;
}