#include <Windows.h>
#include <stdio.h>
#include <string.h>
#include <wchar.h>

static void printHelp();
static bool countDirRecursively(
    const wchar_t *srcPath,
    int& numFiles, int& numDirs,
    double& totalBytes
);

int main() {
    /* Get command line args in Unicode */
    wchar_t *commandLine = GetCommandLineW();
    int argc;
    wchar_t **argv = CommandLineToArgvW(
        commandLine, &argc
    );
    if (argc < 2) {
        LocalFree(argv);
        printHelp();
        return 1;
    }
    DWORD attr = GetFileAttributesW(argv[1]);
    if (attr == INVALID_FILE_ATTRIBUTES) {
        printf("Illegal source directory.\n");
        LocalFree(argv);
        return (-1);
    } else if ((attr & FILE_ATTRIBUTE_DIRECTORY) == 0) {
        printf("Source is not a directory.\n");
        LocalFree(argv);
        return (-1);
    }

    int numFiles = 0;
    int numDirs = 0;
    double totalBytes = 0.;
    if (!countDirRecursively(
        argv[1], 
        numFiles, numDirs, totalBytes
    )) {
        LocalFree(argv);
        return (-1);
    }
    printf(
        "Files: %d, dirs: %d, bytes: %.0lf\n",
        numFiles, numDirs, totalBytes
    );
    return 0;
}

static void printHelp() {
    printf(
        "Count recursively the number of files and\n"
        "subdirectories in a given directory, and their total size.\n"
        "Usage:\n"
        "\trcount directory\n"
    );
}

static bool countDirRecursively(
    const wchar_t *srcPath,
    int& numFiles, int& numDirs,
    double& totalBytes
) {
    DWORD attr = GetFileAttributesW(srcPath);
    if (
        attr == INVALID_FILE_ATTRIBUTES ||
        (attr & FILE_ATTRIBUTE_DIRECTORY) == 0
    ) {
        return false;
    }

    wchar_t srcFile[MAX_PATH];
    wcscpy(srcFile, srcPath);
    size_t len = wcslen(srcFile);
    if (len == 0)
        return false;
    if (srcFile[len-1] != L'\\' && srcFile[len-1] != L'/')
        wcscat(srcFile, L"\\");
    int srcLen = wcslen(srcFile);
    
    wchar_t pattern[MAX_PATH];
    wcscpy(pattern, srcFile);
    wcscat(pattern, L"*");

    numFiles = 0;
    numDirs = 0;
    totalBytes = 0.;

    WIN32_FIND_DATAW findData;
    HANDLE h;
    h = FindFirstFileW(pattern, &findData);
    while (h != INVALID_HANDLE_VALUE) {
        if (
            wcscmp(findData.cFileName, L"..") != 0 &&
            wcscmp(findData.cFileName, L".") != 0
        ) {
            wcscpy(srcFile + srcLen, findData.cFileName);
            printf("%ls\n", srcFile);

            totalBytes += findData.nFileSizeLow;
            // findData.nFileSizeHigh

            if (
                (findData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) != 0
            )
            {
                // Subdirectory
                ++numDirs;
                int nFiles = 0;
                int nDirs = 0;
                double bytes = 0.;
                if (!countDirRecursively(
                    srcFile,
                    nFiles, nDirs, bytes
                )) {
                    printf(
                        "Failed to process subdirectory %ls\n",
                        findData.cFileName
                    );
                    return false;
                }
                numFiles += nFiles;
                numDirs += nDirs;
                totalBytes += bytes;
            } else {
                // Normal file
                ++numFiles;
            }
        }

        if (!FindNextFileW(h, &findData)) 
            break;
    }
    FindClose(h);
        
    return true;
}