#include "misc.h"
#include "resolver.h"

std::random_device dev;
std::mt19937 rng(dev());

std::vector<std::pair<bool, std::string>> cachedLogLines;

void die()
{
    ::ExitProcess(0);
}

void _output(bool verbose, const std::string& out)
{
    if (globalQuietOption) return;

    if (globalCacheLoglines)
    {
        cachedLogLines.push_back(std::make_pair(verbose, out));
        return;
    }

    DWORD written = 0;

    static auto _WriteFile = reinterpret_cast<fn_WriteFile*>(::GetProcAddress(
        GetModuleHandleW(L"kernel32.dll"), "WriteFile"
    ));

    if (globalLogFilePath[0] != L'\0' && globalLogFilePath[0] != L'-')
    {
        static auto _CreateFileW = reinterpret_cast<fn_CreateFileW*>(::GetProcAddress(
            GetModuleHandleW(L"kernel32.dll"), "CreateFileW"
        ));

        HANDLE hFile = _CreateFileW(
            globalLogFilePath,
            FILE_APPEND_DATA,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            NULL,
            OPEN_ALWAYS,
            FILE_ATTRIBUTE_NORMAL,
            NULL
        );

        if (hFile != INVALID_HANDLE_VALUE && hFile != NULL)
        {
            static auto _SetFilePointer = reinterpret_cast<fn_SetFilePointer*>(::GetProcAddress(
                GetModuleHandleW(L"kernel32.dll"), "SetFilePointer"
            ));

            _SetFilePointer(
                hFile,
                0,
                nullptr,
                FILE_END
            );

            _WriteFile(
                hFile,
                out.c_str(),
                static_cast<DWORD>(out.size()),
                &written,
                nullptr
            );

            CloseHandle(hFile);
        }
    }
    else
    {
        _WriteFile(
            GetStdHandle(static_cast<DWORD>(-11) /* STD_OUTPUT_HANDLE */),
            out.c_str(),
            static_cast<DWORD>(out.size()),
            &written,
            nullptr
        );
    }
}

void eraseVerboseLinesInCache()
{
	for (auto it = cachedLogLines.begin(); it != cachedLogLines.end(); )
	{
		if (it->first) it = cachedLogLines.erase(it);
		else ++it;
	}
}

void flushCachedLogLines(bool forceFlushVerbose)
{
    if (!globalCacheLoglines) return;

    globalCacheLoglines = false;

    if (!globalQuietOption)
    {
        for (const auto &line : cachedLogLines)
        {
            if (!line.first) _output(false, line.second);
            else if (line.first && (forceFlushVerbose || globalVerboseOption)) _output(true, line.second);
        }
    }

    cachedLogLines.clear();
}

int getRandomNumber(int a)
{
    return getRandomNumber(0, a - 1);
}

int getRandomNumber(int a, int b)
{
    if (a - b == 0) return a;

    std::uniform_int_distribution<std::mt19937::result_type> dist(a, b);
    return dist(rng);
}

std::wstring adjustPath(const std::wstring& szPath)
{
    auto out = _adjustPath(szPath);
    if (!out.empty()) return out;

    out = _adjustPath(szPath + OBFI(L".exe"));
    if (!out.empty()) return out;

    out = _adjustPath(szPath + OBFI(L".dll"));
    if (!out.empty()) return out;

    out = _adjustPath(OBFI(L"..\\") + szPath);
    if (!out.empty()) return out;

    out = _adjustPath(OBFI(L"..\\") + szPath + OBFI(L".exe"));
    if (!out.empty()) return out;

    out = _adjustPath(OBFI(L"..\\") + szPath + OBFI(L".dll"));
    if (!out.empty()) return out;

    info(OBF(L"[!] Specified file ("), szPath, OBF(L") does not exist in CWD, Windows or Windows\\System32!"));
    die();
    return L"";
}

std::wstring _adjustPath(const std::wstring& szPath)
{
    //
    // Can't RESOLVE GetFileAttributesW and ExpandEnvironmentStringsW as we're unable to
    // #include Resolver here. This would lead to cyclic dependency and wouldn't build.
    //

    DWORD dwAttrib = GetFileAttributesW(szPath.c_str());

    bool exists = (dwAttrib != INVALID_FILE_ATTRIBUTES &&
        !(dwAttrib & FILE_ATTRIBUTE_DIRECTORY));

    if (exists)
    {
        return std::wstring(szPath);
    }
    else
    {
        wchar_t path[1024] = { 0 };

        // Do not alter below ExpandEnvironmentStringsW to RESOLVE as it will introduce
        // infinite recursion!
        ExpandEnvironmentStringsW(OBF_WSTR(L"%SystemRoot%").c_str(), path, _countof(path));

        auto newPath = std::wstring(path) + OBF_WSTR(L"\\System32\\") + szPath;

        dwAttrib = GetFileAttributesW(newPath.c_str());
        exists = (dwAttrib != INVALID_FILE_ATTRIBUTES &&
            !(dwAttrib & FILE_ATTRIBUTE_DIRECTORY));

        if (!exists)
        {
            return {};
        }

        return newPath;
    }
}

std::string adjustPathA(const std::string& szPath)
{
    static std::map<std::string, std::string> cachedAdjustedPaths;
    if (cachedAdjustedPaths.count(szPath) != 0)
    {
        return cachedAdjustedPaths[szPath];
    }

    std::wstring b(szPath.begin(), szPath.end());
    auto a = adjustPath(b);

    cachedAdjustedPaths[szPath] = std::string(a.begin(), a.end());
    return cachedAdjustedPaths[szPath];
}

std::wstring ltrim(const std::wstring& str, const std::wstring& chars)
{
    auto str2 = str;
    str2.erase(0, str.find_first_not_of(chars));
    return str2;
}

std::wstring rtrim(const std::wstring& str, const std::wstring& chars)
{
    auto str2 = str;
    str2.erase(str.find_last_not_of(chars) + 1);
    return str2;
}

std::wstring trim(const std::wstring& str, const std::wstring& chars)
{
    return ltrim(rtrim(str, chars), chars);
}