/**
 * @file 
 * @author George Mount
 *
 * @brief Mock Windows Registry functions
 *
 * Copyright (c) 2009 NVIDIA Corporation. All rights reserved.
 *
 * NVIDIA Corporation and its licensors retain all intellectual property and 
 * proprietary rights in and to this software, related documentation and any 
 * modifications thereto. Any use, reproduction, disclosure or distribution of 
 * this software and related documentation without express license agreement 
 * from NVIDIA Corporation is strictly prohibited.
 */

#include "MockWinApi.h"
#include <boost/foreach.hpp>
#include <boost/assign.hpp>

using namespace Nvidia::Logging::Test;
using namespace boost;
using namespace boost::assign;
using namespace std;

RegistryEntry HKLM_Entry = {HKEY_LOCAL_MACHINE, RegistryEntry::ValuesMap(), RegistryEntry::SubKeysMap()};
RegistryEntry HKCU_Entry = {HKEY_LOCAL_MACHINE, RegistryEntry::ValuesMap(), RegistryEntry::SubKeysMap()};

MockWinApi::RegistryMap MockWinApi::m_registry = map_list_of
    (HKEY_LOCAL_MACHINE, shared_ptr<RegistryEntry>(new RegistryEntry(HKLM_Entry)))
    (HKEY_CURRENT_USER, shared_ptr<RegistryEntry>(new RegistryEntry(HKCU_Entry)));

long MockWinApi::m_returnCode = ERROR_SUCCESS;
DWORD MockWinApi::m_createFileFlags = 0;

vector<tstring> MockWinApi::SubKeys(const tstring &subkey)
{
    vector<tstring> tokens;

    if (subkey.size() > 0)
    {
        const char delimiter = '\\';
        tstring::size_type lastPos = subkey.find_first_not_of(delimiter, 0);
        tstring::size_type pos     = subkey.find_first_of(delimiter, lastPos);

        while (pos != tstring::npos || lastPos != tstring::npos)
        {
            tokens.push_back(subkey.substr(lastPos, pos - lastPos));
            lastPos = subkey.find_first_not_of(delimiter, pos);
            pos = subkey.find_first_of(delimiter, lastPos);
        }
    }
    return tokens;
}

size_t MockWinApi::GetClosestKey(boost::shared_ptr<RegistryEntry> &key, const vector<tstring> &subkeys)
{
    size_t existing = 0;
    while (existing < subkeys.size() && key->SubKeys.find(subkeys[existing]) != key->SubKeys.end())
    {
        key = key->SubKeys[subkeys[existing]];
        existing++;
    }
    return existing;
}

LSTATUS MockWinApi::RegCreateKeyEx (
    __in HKEY hKey,
    __in LPCTSTR lpSubKey,
    __reserved DWORD Reserved,
    __in_opt LPTSTR lpClass,
    __in DWORD dwOptions,
    __in REGSAM samDesired,
    __in_opt CONST LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    __out PHKEY phkResult,
    __out_opt LPDWORD lpdwDisposition
    )
{
    UNREFERENCED_PARAMETER(lpdwDisposition);
    UNREFERENCED_PARAMETER(lpSecurityAttributes);
    UNREFERENCED_PARAMETER(samDesired);
    UNREFERENCED_PARAMETER(dwOptions);
    UNREFERENCED_PARAMETER(lpClass);
    UNREFERENCED_PARAMETER(Reserved);
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        vector<tstring> subkeys = SubKeys(lpSubKey);
        shared_ptr<RegistryEntry> parent = iter->second;
        size_t numKeys = GetClosestKey(parent, subkeys);

        while (numKeys < subkeys.size())
        {
            shared_ptr<RegistryEntry> newSubKey(new RegistryEntry);
            HKEY newKey = (HKEY) newSubKey.get();
            newSubKey->Key = newKey;
            parent->SubKeys[subkeys[numKeys]] = newSubKey;
            m_registry[newKey] = newSubKey;
            parent = newSubKey;
            numKeys++;
        }
        *phkResult = parent->Key;
    }
    return m_returnCode;
}

LSTATUS MockWinApi::RegOpenKeyEx (
    __in HKEY hKey,
    __in_opt LPCTSTR lpSubKey,
    __in_opt DWORD ulOptions,
    __in REGSAM samDesired,
    __out PHKEY phkResult
    )
{
    UNREFERENCED_PARAMETER(samDesired);
    UNREFERENCED_PARAMETER(ulOptions);
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        shared_ptr<RegistryEntry> key = iter->second;
        if (lpSubKey != NULL)
        {
            vector<tstring> subkeys = SubKeys(lpSubKey);
            size_t numKeys = GetClosestKey(key, subkeys);
            if (numKeys < subkeys.size())
            {
                return ERROR_CANTOPEN;
            }
        }
        *phkResult = key->Key;
    }
    return m_returnCode;
}

LSTATUS MockWinApi::SHDeleteKey (
    __in HKEY hKey,
    __in LPCTSTR lpSubKey
    )
{
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        vector<tstring> subkeys = SubKeys(lpSubKey);
        tstring lastKey = subkeys[subkeys.size() - 1];
        subkeys.erase(subkeys.begin() + (subkeys.size() - 1));

        shared_ptr<RegistryEntry> parent = iter->second;
        size_t numKeys = GetClosestKey(parent, subkeys);

        if (numKeys < subkeys.size())
        {
            return ERROR_CANTOPEN;
        }
        parent->SubKeys.erase(lastKey);
    }
    return m_returnCode;
}

LSTATUS MockWinApi::RegDeleteValue (
    __in HKEY hKey,
    __in_opt LPCTSTR lpValueName
    )
{
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        shared_ptr<RegistryEntry> parent = iter->second;
        parent->Values.erase(lpValueName);
    }
    return m_returnCode;
}

LSTATUS MockWinApi::RegSetValueEx (
    __in HKEY hKey,
    __in_opt LPCTSTR lpValueName,
    __reserved DWORD Reserved,
    __in DWORD dwType,
    __in_bcount_opt(cbData) CONST BYTE* lpData,
    __in DWORD cbData
    )
{
    UNREFERENCED_PARAMETER(Reserved);
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        shared_ptr<RegistryEntry> parent = iter->second;
        boost::shared_array<char> buf(new char[cbData]);
        memcpy(buf.get(), lpData, cbData);
        RegistryValue value = {dwType, cbData, buf};
        parent->Values[lpValueName] = value;
    }
    return m_returnCode;
}

LSTATUS MockWinApi::RegEnumValue (
    __in HKEY hKey,
    __in DWORD dwIndex,
    __out_ecount_part_opt(*lpcchValueName, *lpcchValueName + 1) LPTSTR lpValueName,
    __inout LPDWORD lpcchValueName,
    __reserved LPDWORD lpReserved,
    __out_opt LPDWORD lpType,
    __out_bcount_part_opt(*lpcbData, *lpcbData) __out_data_source(REGISTRY) LPBYTE lpData,
    __inout_opt LPDWORD lpcbData
    )
{
    UNREFERENCED_PARAMETER(lpReserved);
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        shared_ptr<RegistryEntry> entry = iter->second;

        if (dwIndex >= entry->Values.size())
        {
            return ERROR_NO_MORE_ITEMS;
        }

        RegistryEntry::ValuesMap::iterator valueIter = entry->Values.begin();
        advance(valueIter, dwIndex);
        const tstring &valueName = valueIter->first;
        const RegistryValue &value = valueIter->second;

        size_t valueBufSize = *lpcchValueName;
        *lpcchValueName = (DWORD)valueName.size();

        if (lpType != NULL)
        {
            *lpType = value.Type;
        }
        if (lpValueName != NULL)
        {
            if (valueBufSize <= valueName.size())
            {
                return ERROR_MORE_DATA;
            }
            memcpy(lpValueName, valueName.c_str(), (valueName.size() + 1) * sizeof(TCHAR));
        }
        size_t dataSize = 0;
        if (lpcbData != NULL)
        {
            dataSize = *lpcbData;
            *lpcbData = (DWORD) value.BufLength;
        }
        if (lpData != NULL)
        {
            if (dataSize < value.BufLength)
            {
                return ERROR_MORE_DATA;
            }
            memcpy(lpData, value.Buf.get(), value.BufLength);
        }
    }
    return m_returnCode;
}

LSTATUS MockWinApi::RegQueryValueEx (
    __in HKEY hKey,
    __in_opt LPCTSTR lpValueName,
    __reserved LPDWORD lpReserved,
    __out_opt LPDWORD lpType,
    __out_bcount_part_opt(*lpcbData, *lpcbData) __out_data_source(REGISTRY) LPBYTE lpData,
    __inout_opt LPDWORD lpcbData
    )
{
    UNREFERENCED_PARAMETER(lpReserved);
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        shared_ptr<RegistryEntry> parent = iter->second;
        if (lpValueName == NULL)
        {
            lpValueName = _T("");
        }
        if (parent->Values.find(lpValueName) == parent->Values.end())
        {
            return ERROR_FILE_NOT_FOUND;
        }
        const RegistryValue &value = parent->Values[lpValueName];
        if (lpType != NULL)
        {
            *lpType = value.Type;
        }
        size_t bufSize = 0;
        if (lpcbData != NULL)
        {
            bufSize = *lpcbData;
            *lpcbData = (DWORD) value.BufLength;
        }
        if (lpData != NULL)
        {
            if (bufSize < value.BufLength)
            {
                return ERROR_MORE_DATA;
            }
            memcpy(lpData, value.Buf.get(), value.BufLength);
        }
    }
    return m_returnCode;
}

LSTATUS MockWinApi::RegEnumKeyEx (
    __in HKEY hKey,
    __in DWORD dwIndex,
    __out_ecount_part_opt(*lpcchName, *lpcchName + 1) LPTSTR lpName,
    __inout LPDWORD lpcchName,
    __reserved LPDWORD lpReserved,
    __out_ecount_part_opt(*lpcchClass,*lpcchClass + 1) LPTSTR lpClass,
    __inout_opt LPDWORD lpcchClass,
    __out_opt PFILETIME lpftLastWriteTime
    )
{
    UNREFERENCED_PARAMETER(lpftLastWriteTime);
    UNREFERENCED_PARAMETER(lpcchClass);
    UNREFERENCED_PARAMETER(lpClass);
    UNREFERENCED_PARAMETER(lpReserved);
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        shared_ptr<RegistryEntry> entry = iter->second;

        if (dwIndex >= entry->SubKeys.size())
        {
            return ERROR_NO_MORE_ITEMS;
        }
        RegistryEntry::SubKeysMap::iterator keysIter = entry->SubKeys.begin();
        advance(keysIter, dwIndex);
        const tstring &keyName = keysIter->first;
        shared_ptr<RegistryEntry> key = keysIter->second;

        size_t nameBufSize = *lpcchName;
        *lpcchName = (DWORD)keyName.size();
        if (lpName != NULL)
        {
            if (nameBufSize <= keyName.size())
            {
                return ERROR_MORE_DATA;
            }
            memcpy(lpName, keyName.c_str(), (keyName.size() + 1) * sizeof(TCHAR));
        }
    }
    return m_returnCode;
}

LSTATUS MockWinApi::RegQueryInfoKey (
    __in HKEY hKey,
    __out_ecount_part_opt(*lpcchClass, *lpcchClass + 1) LPTSTR lpClass,
    __inout_opt LPDWORD lpcchClass,
    __reserved LPDWORD lpReserved,
    __out_opt LPDWORD lpcSubKeys,
    __out_opt LPDWORD lpcbMaxSubKeyLen,
    __out_opt LPDWORD lpcbMaxClassLen,
    __out_opt LPDWORD lpcValues,
    __out_opt LPDWORD lpcbMaxValueNameLen,
    __out_opt LPDWORD lpcbMaxValueLen,
    __out_opt LPDWORD lpcbSecurityDescriptor,
    __out_opt PFILETIME lpftLastWriteTime
    )
{
    UNREFERENCED_PARAMETER(lpClass);
    UNREFERENCED_PARAMETER(lpcchClass);
    UNREFERENCED_PARAMETER(lpReserved);
    UNREFERENCED_PARAMETER(lpcbMaxClassLen);
    UNREFERENCED_PARAMETER(lpcbSecurityDescriptor);
    UNREFERENCED_PARAMETER(lpftLastWriteTime);
    if (m_returnCode == ERROR_SUCCESS)
    {
        RegistryMap::iterator iter = m_registry.find(hKey);

        if (iter == m_registry.end())
        {
            return ERROR_BADKEY;
        }

        shared_ptr<RegistryEntry> entry = iter->second;
        if (lpcSubKeys != NULL)
        {
            *lpcSubKeys = (DWORD)entry->SubKeys.size();
        }
        if (lpcbMaxSubKeyLen != NULL)
        {
            size_t maxLength = 0;
            BOOST_FOREACH (const RegistryEntry::SubKeysMap::value_type &subKey, entry->SubKeys)
            {
                maxLength = max(maxLength, subKey.first.size());
            }
            *lpcbMaxSubKeyLen = (DWORD)maxLength;
        }
        if (lpcValues != NULL)
        {
            *lpcValues = (DWORD)entry->Values.size();
        }
        if (lpcbMaxValueNameLen != NULL)
        {
            size_t maxLength = 0;
            BOOST_FOREACH (const RegistryEntry::ValuesMap::value_type &value, entry->Values)
            {
                maxLength = max(maxLength, value.first.size());
            }
            *lpcbMaxValueNameLen = (DWORD)maxLength;
        }
        if (lpcbMaxValueLen != NULL)
        {
            size_t maxLength = 0;
            BOOST_FOREACH (const RegistryEntry::ValuesMap::value_type &value, entry->Values)
            {
                maxLength = max(maxLength, value.second.BufLength);
            }
            *lpcbMaxValueLen = (DWORD)maxLength;
        }
    }
    return m_returnCode;
}

LSTATUS MockWinApi::RegCloseKey (
    __in HKEY hKey
    )
{
    UNREFERENCED_PARAMETER(hKey);
    return m_returnCode;
}

LSTATUS MockWinApi::RegOpenCurrentUser(
  __in REGSAM samDesired,
  __out PHKEY phkResult
  )
{
    UNREFERENCED_PARAMETER(samDesired);
    *phkResult = HKEY_CURRENT_USER;
    return m_returnCode;
}

HANDLE MockWinApi::CreateFile(
                   __in     LPCTSTR lpFileName,
                   __in     DWORD dwDesiredAccess,
                   __in     DWORD dwShareMode,
                   __in_opt LPSECURITY_ATTRIBUTES lpSecurityAttributes,
                   __in     DWORD dwCreationDisposition,
                   __in     DWORD dwFlagsAndAttributes,
                   __in_opt HANDLE hTemplateFile
                   )
{
    m_createFileFlags = dwFlagsAndAttributes;
    return ::CreateFile(lpFileName, dwDesiredAccess, dwShareMode, 
        lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes,
        hTemplateFile);
}