#include <windows.h>
#include <shlwapi.h>

#include <assert.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "hooklib/path.h"
#include "hooklib/reg.h"

#include "platform/vfs.h"

#include "util/dprintf.h"

static void vfs_fixup_path(wchar_t *path, size_t max_count);
static HRESULT vfs_mkdir_rec(const wchar_t *path);
static HRESULT vfs_path_hook(const wchar_t *src, wchar_t *dest, size_t *count);
static HRESULT vfs_path_hook_nthome(
        const wchar_t *src,
        wchar_t *dest,
        size_t *count);
static HRESULT vfs_path_hook_option(
        const wchar_t *src,
        wchar_t *dest,
        size_t *count);


static struct vfs_config vfs_config;

HRESULT vfs_hook_init(const struct vfs_config *config)
{
    HRESULT hr;

    assert(config != NULL);

    if (!config->enable) {
        return S_FALSE;
    }

    memcpy(&vfs_config, config, sizeof(*config));

    if (config->path[0] == L'\0') {
        dprintf("VFS: !FATAL! Path not set for VFS!\n");
        return E_FAIL;
    }

    vfs_fixup_path(vfs_config.path, _countof(vfs_config.path));
    vfs_fixup_path(vfs_config.d, _countof(vfs_config.d));
    vfs_fixup_path(vfs_config.e, _countof(vfs_config.e));
    vfs_fixup_path(vfs_config.f, _countof(vfs_config.f));
    vfs_fixup_path(vfs_config.g, _countof(vfs_config.g));
    vfs_fixup_path(vfs_config.h, _countof(vfs_config.h));
    vfs_fixup_path(vfs_config.j, _countof(vfs_config.j));

    hr = vfs_mkdir_rec(vfs_config.path);

    if (FAILED(hr)) {
        dprintf("Vfs: Failed to create dir %S: %x\n",
                config->path,
                (int) hr);
        return E_FAIL;
    }

    hr = vfs_mkdir_rec(vfs_config.d);

    if (FAILED(hr)) {
        dprintf("Vfs: Failed to create dir %S: %x\n",
                vfs_config.d,
                (int) hr);
        return E_FAIL;
    }

    hr = vfs_mkdir_rec(vfs_config.e);

    if (FAILED(hr)) {
        dprintf("Vfs: Failed to create dir %S: %x\n",
                vfs_config.e,
                (int) hr);
        return E_FAIL;
    }
    
    hr = vfs_mkdir_rec(vfs_config.f);

    if (FAILED(hr)) {
        dprintf("Vfs: Failed to create dir %S: %x\n",
                vfs_config.f,
                (int) hr);
        return E_FAIL;
    }

    hr = vfs_mkdir_rec(vfs_config.g);

    if (FAILED(hr)) {
        dprintf("Vfs: Failed to create dir %S: %x\n",
                vfs_config.g,
                (int) hr);
        return E_FAIL;
    }
    
    hr = vfs_mkdir_rec(vfs_config.h);

    if (FAILED(hr)) {
        dprintf("Vfs: Failed to create dir %S: %x\n",
                vfs_config.h,
                (int) hr);
        return E_FAIL;
    }
    
    hr = vfs_mkdir_rec(vfs_config.j);

    if (FAILED(hr)) {
        dprintf("Vfs: Failed to create dir %S: %x\n",
                vfs_config.j,
                (int) hr);
        return E_FAIL;
    }

    /* Not auto-creating option directory as it is normally a read-only mount */

    hr = path_hook_push(vfs_path_hook);

    if (FAILED(hr)) {
        return hr;
    }

    return S_OK;
}

static void vfs_fixup_path(wchar_t *path, size_t max_count)
{
    size_t count;
    wchar_t abspath[MAX_PATH];

    assert(path != NULL);
    /* Requirement for PathIsRelativeW */
    assert(max_count <= MAX_PATH);

    if (PathIsRelativeW(path)) {
        count = GetFullPathNameW(path, _countof(abspath), abspath, NULL);

        /* GetFullPathName's length return value is tricky, because it includes
           the NUL terminator on failure, but doesn't on success.
           Check if it fits the temp buf (else it's a failure and includes NUL
           anyway), then if it fits the target buf, NUL included. */
        if (count == 0 || count > _countof(abspath) || count >= max_count) {
            goto fail;
        }

        wcscpy_s(path, max_count, abspath);
    } else {
        count = wcslen(path);
    }

    if (path_is_separator_w(path[count - 1])) {
        return;
    }

    if (count + 2 > max_count) {
        goto fail;
    }

    path[count + 0] = L'\\';
    path[count + 1] = L'\0';
    return;

fail:
    dprintf("Vfs: FATAL: Path too long: %S\n", path);
    abort();
}

static HRESULT vfs_mkdir_rec(const wchar_t *path)
{
    wchar_t *copy;
    wchar_t *pos;
    wchar_t wc;
    HRESULT hr;
    DWORD attr;
    BOOL ok;

    assert(path != NULL);

    copy = _wcsdup(path);

    if (copy == NULL) {
        hr = HRESULT_FROM_WIN32(ERROR_OUTOFMEMORY);

        goto end;
    }

    pos = copy;

    do {
        wc = *pos;

        if (wc == L'\0' || wc == L'/' || wc == L'\\') {
            *pos = L'\0';
            attr = GetFileAttributesW(copy);

            if (attr == INVALID_FILE_ATTRIBUTES) {
                ok = CreateDirectoryW(copy, NULL);

                if (!ok) {
                    hr = HRESULT_FROM_WIN32(GetLastError());

                    goto end;
                }
            }

            *pos = wc;
        }

        pos++;
    } while (wc != L'\0');

    hr = S_OK;

end:
    free(copy);

    return hr;
}

static HRESULT vfs_path_hook(const wchar_t *src, wchar_t *dest, size_t *count)
{
    const wchar_t *redir;
    size_t required;
    size_t redir_len;
    HRESULT hr;

    assert(src != NULL);
    assert(count != NULL);

    if (src[0] == L'\0' || src[1] != L':' || !path_is_separator_w(src[2])) {
        return S_FALSE;
    }

    switch (src[0]) {
    case L'd':
    case L'D':
        redir = vfs_config.d;

        break;
    
    case L'e':
    case L'E':
        redir = vfs_config.e;
        break;

    case L'f':
    case L'F':
        redir = vfs_config.f;
        break;

    case L'g':
    case L'G':
        redir = vfs_config.g;
        break;

    case L'h':
    case L'H':
        redir = vfs_config.h;
        break;

    case L'j':
    case L'J':
        redir = vfs_config.j;
        break;

    default:
        return S_FALSE;
    }

    /* Cut off <prefix>\, replace with redir path, count NUL terminator */

    redir_len = wcslen(redir);
    required = wcslen(src) - 3 + redir_len + 1;

    if (dest != NULL) {
        if (required > *count) {
            return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
        }

        wcscpy_s(dest, *count, redir);
        wcscpy_s(dest + redir_len, *count - redir_len, src + 3);
    }

    *count = required;
    
#if 0
    dprintf("VFS: Redirecting %ls to %ls\n", src, dest);
#endif

    return S_OK;
}