#include <string>
#include <vector>
#include <Windows.h>

using namespace std;

static vector<wstring> FindFiles(const wstring& searchPath, DWORD fileAttributeMask, DWORD fileAttributeNotMask)
{
    vector<wstring> result;
    WIN32_FIND_DATA findData;
    DWORD lastError;
    bool findNextFileFailed = false;

    auto searchOperation = (fileAttributeMask & FILE_ATTRIBUTE_DIRECTORY) ? FindExSearchLimitToDirectories : FindExSearchNameMatch;
    auto searchHandle = FindFirstFileEx(searchPath.c_str(), FindExInfoStandard, &findData, searchOperation, nullptr, 0);
    if (searchHandle == INVALID_HANDLE_VALUE) return result;

    do
    {
        if (_wcsicmp(findData.cFileName, L".") != 0 && _wcsicmp(findData.cFileName, L"..") != 0)
        {
            if ((findData.dwFileAttributes & fileAttributeMask) && (findData.dwFileAttributes & fileAttributeNotMask) == 0)
            {
                result.push_back(findData.cFileName);

                if (findNextFileFailed)
                {
                    // If this gets hit, that means FindNextFile failed at least once
                    // before succeeding and returning a file
                    __debugbreak();
                }
            }
        }

        FindNextFile(searchHandle, &findData);

        lastError = GetLastError();
        if (lastError != ERROR_SUCCESS && lastError != ERROR_NO_MORE_FILES)
        {
            findNextFileFailed = true;
            OutputDebugString((L"Something went bad while enumerating \"" + searchPath + L"\": " + to_wstring(GetLastError()) + L"\r\n").c_str());
        }
    }
    while (lastError != ERROR_NO_MORE_FILES);

    FindClose(searchHandle);
    SetLastError(ERROR_SUCCESS);

    return result;
}

vector<wstring> GetFilesInDirectory(wstring path, const wstring& searchPattern, bool recursive)
{
    vector<wstring> result;

    if (path[path.length() - 1] != L'\\')
    {
        path += L'\\';
    }

    // Find folders
    if (recursive)
    {
        for (const auto& directory : FindFiles(path + L"*.*", FILE_ATTRIBUTE_DIRECTORY, 0))
        {
            for (auto& file : GetFilesInDirectory(path + directory, searchPattern, recursive))
            {
                result.emplace_back(std::move(file));
            }
        }
    }

    // Find files
    for (const auto& fileName : FindFiles(path + searchPattern, 0xFFFFFFFF, FILE_ATTRIBUTE_DIRECTORY))
    {
        result.emplace_back(path + fileName);
    }

    return result;
}

int main()
{
    auto files = GetFilesInDirectory(L"C:\\Windows\\System32", L"*.*", true);

    return 0;
}