#include <windows.h>
#include <winnls.h>
#include <shobjidl.h>
#include <objbase.h>
#include <Shlwapi.h>
#include <objidl.h>
#include <shlguid.h>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <thread>
#include "hid_handler.h"
#include "process_watch.h"
#include "utils.h"
#include "ui.h"
#include "localization.h"
#include "user_data_manager.h"

// Is this process running with elevated rights?
BOOL is_elevated()
{
    BOOL fRet = FALSE;
    HANDLE hToken = NULL;
    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
    {
        TOKEN_ELEVATION Elevation;
        DWORD cbSize = sizeof(TOKEN_ELEVATION);
        if (GetTokenInformation(hToken, TokenElevation, &Elevation, sizeof(Elevation), &cbSize))
        {
            fRet = Elevation.TokenIsElevated;
        }
    }
    if (hToken)
    {
        CloseHandle(hToken);
    }
    return fRet;
}

/// Create a shortcut.
static HRESULT create_link(LPCWSTR working_dir,
    LPCWSTR exe_path,
    LPCWSTR shorcut_path,
    LPCWSTR description)
{
    HRESULT hres;
    IShellLink* psl;
    hres = CoCreateInstance(CLSID_ShellLink, NULL, CLSCTX_INPROC_SERVER, IID_IShellLink, (LPVOID*)&psl);
    if (SUCCEEDED(hres))
    {
        IPersistFile* ppf;
        psl->SetPath(exe_path);
        psl->SetDescription(description);
        psl->SetWorkingDirectory(working_dir);
        psl->SetArguments(L"-startupLaunch");
        hres = psl->QueryInterface(IID_IPersistFile, (LPVOID*)&ppf);
        if (SUCCEEDED(hres))
        {
            hres = ppf->Save(shorcut_path, TRUE);
            ppf->Release();
        }
        psl->Release();
    }
    return hres;
}

bool get_reg_value(const std::wstring& scheme, std::wstring& outValue)
{
    HKEY hKey;
    std::wstring commandKeyPath = L"Software\\Classes\\" + scheme + L"\\shell\\open\\command";

    if (RegOpenKeyExW(HKEY_CURRENT_USER, commandKeyPath.c_str(), 0, KEY_READ, &hKey) != ERROR_SUCCESS)
    {
        return false;
    }

    wchar_t buffer[1024];
    DWORD bufferSize = sizeof(buffer);
    DWORD type = 0;

    LONG result = RegQueryValueExW(hKey, NULL, nullptr, &type, (LPBYTE)buffer, &bufferSize);
    RegCloseKey(hKey);

    if (result != ERROR_SUCCESS || type != REG_SZ)
    {
        return false;
    }

    outValue = buffer;
    return true;
}

void register_uri(const std::wstring& scheme, const std::wstring& exePath)
{
    std::wstring expectedCommand = L"\"" + exePath + L"\" \"%1\"";

    std::wstring existingCommand;
    if (get_reg_value(scheme, existingCommand))
    {
        // Entry already exists. No point in recreating
        if (existingCommand == expectedCommand)
        {
            return;
        }
    }

    HKEY hKey;
    std::wstring baseKeyPath = L"Software\\Classes\\" + scheme;

    if (RegCreateKeyExW(HKEY_CURRENT_USER, baseKeyPath.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE, NULL, &hKey, NULL) != ERROR_SUCCESS)
    {
        return;
    }

    std::wstring description = L"URL:" + scheme + L" Protocol";
    RegSetValueExW(hKey, NULL, 0, REG_SZ, (BYTE*)description.c_str(), static_cast<DWORD>((description.size() + 1) * sizeof(wchar_t)));
    RegSetValueExW(hKey, L"URL Protocol", 0, REG_SZ, (BYTE*)L"", sizeof(wchar_t));

    HKEY hSubKey;
    std::wstring commandKeyPath = baseKeyPath + L"\\shell\\open\\command";
    if (RegCreateKeyExW(HKEY_CURRENT_USER, commandKeyPath.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE, NULL, &hSubKey, NULL) != ERROR_SUCCESS)
    {
        RegCloseKey(hKey);
        return;
    }

    std::wstring command = L"\"" + exePath + L"\" \"%1\"";
    RegSetValueExW(hSubKey, NULL, 0, REG_SZ, (BYTE*)command.c_str(), static_cast<DWORD>((command.size() + 1) * sizeof(wchar_t)));

    RegCloseKey(hSubKey);
    RegCloseKey(hKey);

    return;
}

///
int main(int argc, char* argv[])
{
    bool startClosed = false;
    bool restartInstance = false;
    bool ignoreMutex = false;
    bool hasResetCamera = false;
    bool closeExistingProcess = false;

    // Register the utility URI
    wchar_t cwd[MAX_PATH];
    GetModuleFileNameW(NULL, cwd, MAX_PATH);
    PathRemoveFileSpecW(cwd);

    std::wstring et_path = std::wstring(cwd) + L"\\BeyondHID.exe";
    register_uri(L"bset", et_path);

    // Get current executable path
    char exePath[MAX_PATH];
    DWORD len = GetModuleFileNameA(NULL, exePath, MAX_PATH);

    if (len != 0)
    {
#ifndef _DEBUG
        // Force working directory to the executable
        char exeDir[MAX_PATH];
        strcpy_s(exeDir, exePath);

        if (PathRemoveFileSpecA(exeDir))
        {
            SetCurrentDirectoryA(exeDir);
        }
#endif
    }

    Localization::SetLanguage(Locales::None);

    hid_handler.log.open("log.txt");
    hid_handler.log.flush();

    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "-startupLaunch") == 0)
        {
            startClosed = true;
        }

        if (std::strcmp(argv[i], "-elevated") == 0)
        {
            ignoreMutex = true;
        }

        if (std::strcmp(argv[i], "-restartCamera") == 0)
        {
            if (is_elevated())
            {
                hid_handler.startup_camera_restart.store(true);
            }

            hasResetCamera = true;
            ignoreMutex = true;
        }

        if (std::string(argv[i]).find("bset://", 0) != std::string::npos)
        {
            closeExistingProcess = true;
            
            std::string arg = argv[i];
            size_t pos = arg.find("t=");
            if (pos != std::string::npos)
            {
                ui.get_runtime_state() -> et_user_token = (arg.substr(pos + 2).c_str());
                hid_handler.log << "Command line arguments parsed token!" << std::endl;
            }
        }
    }

    // If program is already running, don't allow a second instance.
    LPCWSTR named_mutex = L"com_bigscreen_hmd_utility";
    HANDLE program_mutex = CreateMutex(NULL, TRUE, named_mutex);
    if (!ignoreMutex && ERROR_ALREADY_EXISTS == GetLastError())
    {
        if (closeExistingProcess)
        {
            if (!ui.close_existing_process())
            {
                return 0;
            }
        }
        else
        {
            // Restore window minimized / background
            if (!ui.bring_foreground(true))
            {
                ui.create_window(startClosed);
                MessageBox(glfwGetWin32Window(ui.window), LWSTR(StringID::ERROR_UTILITY_ALREADY_RUNNING).c_str(), LWSTR(StringID::POPUP_TITLE_ERROR).c_str(), MB_OK);
            }
        
            return 0;
        }
    }

    // Create window
    ui.create_window(startClosed);

    if (hasResetCamera && !is_elevated())
    {
        MessageBox(glfwGetWin32Window(ui.window), LWSTR(StringID::ERROR_ADMIN_PRIVILEGES_REQUIRED).c_str(), LWSTR(StringID::POPUP_TITLE_ERROR).c_str(), MB_OK);
    }

    if (!IsDebuggerPresent())
    {
        const wchar_t* progDataPath = _wgetenv(L"PROGRAMDATA");

        if (progDataPath)
        {
            std::wstring linkPath = std::wstring(progDataPath) + L"\\Microsoft\\Windows\\Start Menu\\Programs\\StartUp\\BeyondHID.lnk";

            DWORD dwAttrib = GetFileAttributes(linkPath.c_str());
            if (dwAttrib != INVALID_FILE_ATTRIBUTES && !(dwAttrib & FILE_ATTRIBUTE_DIRECTORY))
            {
                // If we're already in the startup folder, no need for elevated priviledges

            }
            else
            {
                // Restart with elevated priviledges.
                if (!is_elevated())
                {
                    ShellExecute(NULL, L"runas", L"BeyondHID.exe", L"-elevated", NULL, SW_SHOWNORMAL);
                    return 0;
                }

                // Save a shortcut into the startup folder so it runs on boot.
                CoInitialize(nullptr);
                char exe_path[MAX_PATH];
                GetModuleFileNameA(GetModuleHandle(NULL), exe_path, MAX_PATH);
                WCHAR exe_wstr[MAX_PATH];
                MultiByteToWideChar(CP_ACP, 0, exe_path, -1, exe_wstr, MAX_PATH);
                PathRemoveFileSpecA(exe_path);
                WCHAR dir_wstr[MAX_PATH];
                MultiByteToWideChar(CP_ACP, 0, exe_path, -1, dir_wstr, MAX_PATH);

                HRESULT hr = create_link(dir_wstr,
                    exe_wstr,
                    linkPath.c_str(), L"Bigscreen");
            }
        }
    }

    // Setup states.
#ifndef _DEBUG
    #ifdef APP_VERSION
        ui.app_version = APP_VERSION;
    #endif
#endif

    ui.alive.store(true);
    hid_handler.ready.store(false);

    // HID thread.
    hid_handler.run();

    // Process watch
    process_watch.run();

    // Wait for hid handler to be setup.
    while (!hid_handler.ready.load())
    {
        Sleep(10); // Don't clog the core.
    }

    // UI (main) thread.
    ui.run(); // Blocks until execution stops.

    // Join and cleanup threads.
    hid_handler.shutdown();
    ui.shutdown();
    utils.shutdown();
    process_watch.shutdown();

    CoUninitialize();

    ReleaseMutex(program_mutex);
    CloseHandle(program_mutex);

#ifdef _WIN32
    WSACleanup();
#endif
    return 0;
}